use pagesift_content::Document;
use pagesift_stat::calib::{Isotonic, Logistic};
use pagesift_stat::{DecodeError, Reader, Writer};
use crate::script::ScriptWeights;
pub const FEATURE_NAMES: [&str; 13] = [
"bias",
"ln_text_weight",
"ln_html_bytes",
"coverage_measured",
"coverage_share",
"page_share",
"root_in_comments",
"ln_comment_ratio",
"barrier_confidence",
"is_article",
"ln_headings",
"ln_paragraphs",
"link_text_share",
];
pub type Features = [f64; FEATURE_NAMES.len()];
#[derive(Debug, Default, Clone, Copy, PartialEq)]
struct Shape {
headings: usize,
paragraphs: usize,
visible_chars: usize,
link_chars: usize,
}
fn shape(md: &str) -> Shape {
let mut s = Shape::default();
let mut in_para = false;
let mut in_code = false;
for line in md.lines() {
let t = line.trim_start();
if t.starts_with("```") {
in_code = !in_code;
in_para = false;
continue;
}
if in_code {
continue;
}
if t.is_empty() {
in_para = false;
continue;
}
let t = if let Some(rest) = t.strip_prefix('#') {
s.headings += 1;
in_para = false;
rest.trim_start_matches('#')
} else {
if !in_para {
s.paragraphs += 1;
in_para = true;
}
t
};
let mut depth = 0usize;
let mut skip_url = false;
let mut prev = '\0';
for c in t.chars() {
match c {
'[' => depth += 1,
']' => depth = depth.saturating_sub(1),
'(' if prev == ']' => skip_url = true,
')' if skip_url => skip_url = false,
c if !c.is_whitespace() && !skip_url => {
s.visible_chars += 1;
if depth > 0 {
s.link_chars += 1;
}
}
_ => {}
}
prev = c;
}
}
s
}
pub fn features(doc: &Document, html_bytes: usize, weights: &ScriptWeights) -> Features {
let sh = shape(&doc.markdown);
let text = doc.plain_text.as_deref().unwrap_or(&doc.markdown);
let tw = weights.text_weight(text);
let cov = &doc.coverage;
let out = cov.out_chars() as f64;
let (measured, share) = match cov.main_chars() {
Some(main) if main > 0 => (1.0, (out / main as f64).min(1.0)),
_ => (0.0, 1.0),
};
let comments = cov.comment_chars().unwrap_or(0) as f64;
[
1.0,
tw.ln_1p(),
(html_bytes as f64).ln_1p(),
measured,
share,
if cov.page_chars() == 0 {
1.0
} else {
(out / cov.page_chars() as f64).min(1.0)
},
f64::from(u8::from(cov.root_in_comments())),
(comments / (out + 1.0)).ln_1p(),
doc.barrier
.as_ref()
.map_or(0.0, |b| f64::from(b.confidence)),
f64::from(u8::from(doc.metadata.is_article)),
(sh.headings as f64).ln_1p(),
(sh.paragraphs as f64).ln_1p(),
if sh.visible_chars == 0 {
0.0
} else {
sh.link_chars as f64 / sh.visible_chars as f64
},
]
}
#[derive(Debug, Clone, PartialEq)]
pub struct QualityModel {
pub scores: Vec<Logistic>,
pub calibration: Isotonic,
}
pub const MODEL_VERSION: u32 = 1;
impl QualityModel {
fn fit_score(set: &[(Features, bool)]) -> Option<Logistic> {
let first = set.first()?.0;
let active: Vec<usize> = (0..FEATURE_NAMES.len())
.filter(|&i| i == 0 || set.iter().any(|(f, _)| f[i] != first[i]))
.collect();
let xs: Vec<Vec<f64>> = set
.iter()
.map(|(f, _)| active.iter().map(|&i| f[i]).collect())
.collect();
let ys: Vec<bool> = set.iter().map(|(_, y)| *y).collect();
let reduced = Logistic::fit(&xs, &ys)?;
let mut w = vec![0.0; FEATURE_NAMES.len()];
for (k, &i) in active.iter().enumerate() {
w[i] = reduced.w[k];
}
Some(Logistic { w })
}
pub fn fit(fit_set: &[(Features, bool)], calib_set: &[(Features, bool)]) -> Option<Self> {
let score = Self::fit_score(fit_set)?;
let s: Vec<f64> = calib_set.iter().map(|(f, _)| score.predict(f)).collect();
let y: Vec<bool> = calib_set.iter().map(|(_, y)| *y).collect();
Some(Self {
calibration: Isotonic::fit(&s, &y),
scores: vec![score],
})
}
pub fn fit_cross(a: &[(Features, bool)], b: &[(Features, bool)]) -> Option<Self> {
let (sa, sb) = (Self::fit_score(a)?, Self::fit_score(b)?);
let oof = |part: &[(Features, bool)], m: &Logistic| -> (Vec<f64>, Vec<bool>) {
part.iter().map(|(f, y)| (m.predict(f), *y)).unzip()
};
let (mut s, mut y) = oof(b, &sa);
let (s_a, y_a) = oof(a, &sb);
s.extend(s_a);
y.extend(y_a);
Some(Self {
calibration: Isotonic::fit(&s, &y),
scores: vec![sa, sb],
})
}
pub fn builtin() -> Option<&'static Self> {
static MODEL: std::sync::LazyLock<Option<QualityModel>> = std::sync::LazyLock::new(|| {
QualityModel::decode(include_bytes!("../data/quality-v1.bin"))
.ok()
.flatten()
});
MODEL.as_ref()
}
pub fn predict(&self, f: &Features) -> f64 {
let sum: f64 = self
.scores
.iter()
.map(|m| self.calibration.predict(m.predict(f)))
.sum();
sum / self.scores.len() as f64
}
pub fn encode(&self) -> Vec<u8> {
let mut w = Writer::new();
w.u32(MODEL_VERSION);
w.u32(FEATURE_NAMES.len() as u32);
w.u32(self.scores.len() as u32);
for m in &self.scores {
m.encode(&mut w);
}
self.calibration.encode(&mut w);
w.into_inner()
}
pub fn decode(buf: &[u8]) -> Result<Option<Self>, DecodeError> {
let mut r = Reader::new(buf);
if r.u32()? != MODEL_VERSION || r.u32()? as usize != FEATURE_NAMES.len() {
return Ok(None);
}
let n = r.u32()? as usize;
let scores = (0..n)
.map(|_| Logistic::decode(&mut r))
.collect::<Result<Vec<_>, _>>()?;
let calibration = Isotonic::decode(&mut r)?;
if scores.is_empty() || scores.iter().any(|m| m.w.len() != FEATURE_NAMES.len()) {
return Err(DecodeError::Invalid("quality model dims"));
}
Ok(Some(Self {
scores,
calibration,
}))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn builtin_model_decodes() {
assert!(QualityModel::builtin().is_some());
}
#[test]
fn markdown_shape_counts_blocks_and_link_text() {
let s = shape(
"# T\n\npara one\nstill one\n\n[link](http://x/y) tail\n\n```\n# not heading\n```\n",
);
assert_eq!((s.headings, s.paragraphs), (1, 2));
assert_eq!(s.visible_chars, 1 + 7 + 8 + 4 + 4);
assert_eq!(s.link_chars, 4);
}
#[test]
fn model_round_trip_and_monotone_output() {
let mut data: Vec<(Features, bool)> = Vec::new();
for i in 0..400 {
let mut f = [0.0; FEATURE_NAMES.len()];
f[0] = 1.0;
f[1] = i as f64 / 40.0;
data.push((f, (i * 7919) % 400 < i));
}
let (a, b) = data.split_at(200);
let m = QualityModel::fit(a, b).unwrap();
let mut lo = data[0].0;
let mut hi = data[0].0;
lo[1] = 0.0;
hi[1] = 10.0;
assert!(m.predict(&lo) <= m.predict(&hi));
assert_eq!(QualityModel::decode(&m.encode()).unwrap(), Some(m));
}
}