hypersteeldb 0.5.3

//! **Head C: the biaffine relation scorer** — the head that makes dimension 2 real.
//!
//! Relation polarity is a property of the *argument side*, not of the predicate token, so it cannot be
//! produced by token tagging: knowing that "develops" is a `REL` span says nothing about which entity
//! acts and which is acted upon. Head C binds arguments explicitly. For a predicted pair `(h, t)` with
//! relation `r` the projector emits `rel/r/+` on `h` and `rel/r/-` on `t`.
//!
//! ## Architecture
//!
//! ```text
//!   encoder hidden [T,H] ──span pooling──▶ span reps [S,H']   H' = 2H (start ⊕ mean)
//!                                              │
//!                          biaffine:  s(h,t)_r = hᵀ W_r t  +  U_r·[h;t]  +  b_r
//!                                              │
//!                                       pair logits [S,S,R+1]
//!                                              │
//!                       TYPE MASK from the spec: for facets (f_h, f_t) only relations
//!                       declared `head=f_h, tail=f_t` are scorable; everything else is −∞
//! ```
//!
//! The type mask is the load-bearing detail. `develops: org → system` means the reversed pairing is
//! **unrepresentable**, not merely unlikely — the paper's guarded fragment enforced inside the model, and
//! the reason Head C's output is linter-clean by construction.

use crate::tagger_data::{pair_mask, LabeledSpan, TaggerExample};
use crate::tagger_train::{hrm_config_from, MultiHeadTagger, TrainConfig};
use crate::vocabulary::VocabularySpace;
use candle_core::{DType, Device, IndexOp, Tensor, D};
use candle_nn::{loss, ops::softmax, AdamW, Linear, Module, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use serde::Serialize;
use std::collections::HashSet;
use tokenizers::Tokenizer;

/// Biaffine pair scorer over pooled span representations.
pub struct BiaffineHead {
    /// `[R1, H', H']` — one bilinear matrix per relation class (including `none`)
    w: Tensor,
    /// concatenation term `[h;t] → R1`
    u: Linear,
    n_rel: usize,
    hp: usize,
}

impl BiaffineHead {
    pub fn new(vb: VarBuilder, hp: usize, n_rel: usize) -> candle_core::Result<Self> {
        let w = vb.get((n_rel, hp, hp), "biaffine_w")?;
        let u = candle_nn::linear(2 * hp, n_rel, vb.pp("biaffine_u"))?;
        Ok(BiaffineHead { w, u, n_rel, hp })
    }

    /// Score every ordered pair: `spans [S,H'] → logits [S,S,R1]`.
    pub fn forward(&self, spans: &Tensor) -> candle_core::Result<Tensor> {
        let (s, hp) = spans.dims2()?;
        debug_assert_eq!(hp, self.hp);
        // bilinear term, one relation class at a time: [S,H'] @ W_r [H',H'] @ [H',S] → [S,S]
        let mut planes: Vec<Tensor> = Vec::with_capacity(self.n_rel);
        let spans_t = spans.t()?.contiguous()?;
        for r in 0..self.n_rel {
            let wr = self.w.i(r)?.contiguous()?;
            let bil = spans.matmul(&wr)?.matmul(&spans_t)?; // [S,S]
            planes.push(bil.unsqueeze(2)?); // [S,S,1]
        }
        let bilinear = Tensor::cat(&planes, 2)?; // [S,S,R1]

        // concatenation term: broadcast [h;t] over all pairs → [S,S,R1]
        let h_rep = spans.unsqueeze(1)?.expand((s, s, hp))?; // head varies along dim0
        let t_rep = spans.unsqueeze(0)?.expand((s, s, hp))?; // tail varies along dim1
        let cat = Tensor::cat(&[h_rep, t_rep], 2)?.reshape((s * s, 2 * hp))?;
        let lin = self.u.forward(&cat)?.reshape((s, s, self.n_rel))?;
        bilinear + lin
    }
}

/// Pool one span from encoder hidden states: `start ⊕ mean` over the tokens overlapping the span, giving
/// a `2H` representation. Endpoint+mean is the standard span encoding; it keeps boundary information that
/// a pure mean discards.
pub fn pool_span(hidden: &Tensor, offsets: &[(usize, usize)], span: &LabeledSpan) -> candle_core::Result<Tensor> {
    let idx: Vec<u32> = offsets
        .iter()
        .enumerate()
        .filter(|(_, (ts, te))| te > ts && *ts < span.end && span.start < *te)
        .map(|(i, _)| i as u32)
        .collect();
    let h = hidden.i(0)?; // [T,H]
    if idx.is_empty() {
        // no token overlaps (truncation) → zero vector of the right width
        let dim = h.dim(1)?;
        let z = Tensor::zeros(dim, h.dtype(), h.device())?;
        return Tensor::cat(&[z.clone(), z], 0);
    }
    let sel = Tensor::from_vec(idx.clone(), idx.len(), h.device())?;
    let toks = h.index_select(&sel, 0)?; // [n,H]
    let start = toks.i(0)?;
    let mean = toks.mean(0)?;
    Tensor::cat(&[start, mean], 0)
}

/// Per-pair target class: index into `[none, rel_0, rel_1, …]`.
fn pair_targets(spec: &VocabularySpace, ex: &TaggerExample) -> Vec<Vec<usize>> {
    let names: Vec<&str> = spec.relation_facets.iter().map(|r| r.name.as_str()).collect();
    let n = ex.spans.len();
    let mut t = vec![vec![0usize; n]; n];
    for r in &ex.relations {
        if let Some(ri) = names.iter().position(|n| *n == r.name) {
            if r.head < n && r.tail < n {
                t[r.head][r.tail] = ri + 1;
            }
        }
    }
    t
}

/// The type mask: `true` where relation class `r` is declarable for a pair of facets. Class 0 (`none`) is
/// always allowed; a declared relation is allowed only for its exact `head`/`tail` facets.
pub fn type_allowed(spec: &VocabularySpace, mask: &HashSet<(String, String, String)>, fh: &str, ft: &str, class: usize) -> bool {
    if class == 0 {
        return true;
    }
    match spec.relation_facets.get(class - 1) {
        Some(r) => mask.contains(&(fh.to_string(), ft.to_string(), r.name.clone())),
        None => false,
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct RelReport {
    pub examples: usize,
    pub pairs: usize,
    pub classes: usize,
    pub first_loss: f64,
    pub last_loss: f64,
    /// accuracy over scored pairs after training
    pub train_acc: f64,
    /// accuracy over only the pairs that carry a real relation (not `none`) — the metric that matters,
    /// since `none` dominates
    pub train_acc_positive: f64,
    /// held-out examples (never trained on)
    pub dev_examples: usize,
    /// accuracy on held-out pairs — the only number that says anything about generalisation. Train
    /// accuracy saturates trivially here: the type mask plus one-relation-per-sentence generation leaves
    /// little ambiguity, so a perfect train score is memorisation, not skill.
    pub dev_acc: f64,
    pub dev_acc_positive: f64,
}

/// Train Head C on top of a (frozen) tuned encoder. The encoder is reloaded read-only from the step-2
/// checkpoint: Head A already learned the span representations, so Head C only needs to learn the pair
/// geometry, which keeps this stage cheap and stable.
pub fn train_relations(
    spec: &VocabularySpace,
    examples: &[TaggerExample],
    cfg: &TrainConfig,
    tagger_dir: &std::path::Path,
    epochs: usize,
) -> Result<(VarMap, RelReport), String> {
    let device = Device::Cpu;
    let bert_cfg: candle_transformers::models::bert::Config = serde_json::from_slice(
        &std::fs::read(cfg.base_dir.join("config.json")).map_err(|e| format!("base config: {e}"))?,
    )
    .map_err(|e| format!("parse base config: {e}"))?;
    let tok = Tokenizer::from_file(&cfg.tokenizer).map_err(|e| format!("tokenizer: {e}"))?;
    let n_a = crate::tagger_data::head_a_labels(spec).len();
    let n_b = crate::tagger_data::head_b_labels().len();

    // frozen encoder from step 2
    let weights = tagger_dir.join("tagger.safetensors");
    let vb_frozen = unsafe {
        VarBuilder::from_mmaped_safetensors(&[weights.clone()], DType::F32, &device)
            .map_err(|e| format!("load {}: {e}", weights.display()))?
    };
    let encoder = MultiHeadTagger::new(vb_frozen, &hrm_config_from(&bert_cfg), n_a, n_b).map_err(|e| format!("build encoder: {e}"))?;

    // trainable biaffine head
    let hp = 2 * bert_cfg.hidden_size;
    let n_rel = spec.relation_facets.len() + 1;
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
    let head = BiaffineHead::new(vb, hp, n_rel).map_err(|e| format!("build head C: {e}"))?;
    let mask = pair_mask(spec);

    // pre-encode: (span reps, targets, facets) per example — the encoder is frozen so this is done once
    struct Prepared {
        spans: Tensor, // [S,H']
        targets: Vec<Vec<usize>>,
        facets: Vec<String>,
    }
    let mut prepared: Vec<Prepared> = Vec::new();
    for ex in examples {
        if ex.spans.len() < 2 {
            continue;
        }
        let enc = match tok.encode(ex.text.as_str(), true) {
            Ok(e) => e,
            Err(_) => continue,
        };
        let n = enc.get_ids().len().min(cfg.max_len);
        let ids = Tensor::from_vec(enc.get_ids()[..n].to_vec(), (1, n), &device).map_err(|e| e.to_string())?;
        let attn = Tensor::from_vec(vec![1u32; n], (1, n), &device).map_err(|e| e.to_string())?;
        let hidden = encoder.hidden(&ids, &attn, false).map_err(|e| format!("encode: {e}"))?;
        let offsets: Vec<(usize, usize)> = enc.get_offsets()[..n].to_vec();
        let reps: Vec<Tensor> = ex
            .spans
            .iter()
            .map(|sp| pool_span(&hidden, &offsets, sp))
            .collect::<candle_core::Result<Vec<_>>>()
            .map_err(|e| format!("pool: {e}"))?;
        let spans = Tensor::stack(&reps, 0).map_err(|e| e.to_string())?.detach();
        prepared.push(Prepared { spans, targets: pair_targets(spec, ex), facets: ex.spans.iter().map(|s| s.facet.clone()).collect() });
    }
    if prepared.len() < 5 {
        return Err("too few examples with >=2 spans to split train/dev".into());
    }
    // deterministic 80/20 holdout by stride so every relation/case is represented in both halves
    let mut dev: Vec<Prepared> = Vec::new();
    let mut train_set: Vec<Prepared> = Vec::new();
    for (i, p) in prepared.into_iter().enumerate() {
        if i % 5 == 4 {
            dev.push(p);
        } else {
            train_set.push(p);
        }
    }
    let prepared = train_set;

    let mut opt = AdamW::new(varmap.all_vars(), ParamsAdamW { lr: cfg.lr, ..Default::default() })
        .map_err(|e| format!("optimizer: {e}"))?;
    let (mut first_loss, mut last_loss) = (f64::NAN, f64::NAN);
    let mut total_pairs = 0usize;

    for epoch in 1..=epochs {
        let mut sum = 0.0f64;
        let mut steps = 0usize;
        for p in &prepared {
            let logits = head.forward(&p.spans).map_err(|e| format!("head C forward: {e}"))?;
            let s = p.facets.len();
            // gather type-allowed pairs (excluding the diagonal: a span never relates to itself)
            let mut rows: Vec<Tensor> = Vec::new();
            let mut tgts: Vec<u32> = Vec::new();
            for i in 0..s {
                for j in 0..s {
                    if i == j {
                        continue;
                    }
                    let cls = p.targets[i][j];
                    // mask: disallowed classes get -inf so they are unrepresentable
                    let row = logits.i((i, j)).map_err(|e| e.to_string())?; // [R1]
                    let allow: Vec<f32> = (0..n_rel)
                        .map(|c| if type_allowed(spec, &mask, &p.facets[i], &p.facets[j], c) { 0.0 } else { f32::NEG_INFINITY })
                        .collect();
                    let allow = Tensor::from_vec(allow, n_rel, &device).map_err(|e| e.to_string())?;
                    rows.push((row + allow).map_err(|e| e.to_string())?.unsqueeze(0).map_err(|e| e.to_string())?);
                    tgts.push(cls as u32);
                }
            }
            if rows.is_empty() {
                continue;
            }
            let batch = Tensor::cat(&rows, 0).map_err(|e| e.to_string())?;
            let tgt = Tensor::from_vec(tgts.clone(), tgts.len(), &device).map_err(|e| e.to_string())?;
            let l = loss::cross_entropy(&batch, &tgt).map_err(|e| format!("ce: {e}"))?;
            opt.backward_step(&l).map_err(|e| format!("backward: {e}"))?;
            sum += l.to_scalar::<f32>().map_err(|e| e.to_string())? as f64;
            steps += 1;
            if epoch == 1 {
                total_pairs += tgts.len();
            }
        }
        let avg = sum / steps.max(1) as f64;
        if epoch == 1 {
            first_loss = avg;
        }
        last_loss = avg;
    }

    // accuracy over a set
    let score = |set: &[Prepared]| -> Result<(f64, f64), String> {
        let (mut ok, mut n, mut ok_pos, mut n_pos) = (0usize, 0usize, 0usize, 0usize);
        for p in set {
            let logits = head.forward(&p.spans).map_err(|e| e.to_string())?;
            let s = p.facets.len();
            for i in 0..s {
                for j in 0..s {
                    if i == j {
                        continue;
                    }
                    let row = logits.i((i, j)).map_err(|e| e.to_string())?;
                    let allow: Vec<f32> = (0..n_rel)
                        .map(|c| if type_allowed(spec, &mask, &p.facets[i], &p.facets[j], c) { 0.0 } else { f32::NEG_INFINITY })
                        .collect();
                    let allow = Tensor::from_vec(allow, n_rel, &device).map_err(|e| e.to_string())?;
                    let pred = softmax(&(row + allow).map_err(|e| e.to_string())?, D::Minus1)
                        .and_then(|t| t.argmax(D::Minus1))
                        .and_then(|t| t.to_scalar::<u32>())
                        .map_err(|e| e.to_string())? as usize;
                    let want = p.targets[i][j];
                    n += 1;
                    if pred == want {
                        ok += 1;
                    }
                    if want != 0 {
                        n_pos += 1;
                        if pred == want {
                            ok_pos += 1;
                        }
                    }
                }
            }
        }
        Ok((ok as f64 / n.max(1) as f64, ok_pos as f64 / n_pos.max(1) as f64))
    };
    let (dev_acc, dev_acc_pos) = score(&dev)?;

    // accuracy
    let (mut ok, mut n, mut ok_pos, mut n_pos) = (0usize, 0usize, 0usize, 0usize);
    for p in &prepared {
        let logits = head.forward(&p.spans).map_err(|e| e.to_string())?;
        let s = p.facets.len();
        for i in 0..s {
            for j in 0..s {
                if i == j {
                    continue;
                }
                let row = logits.i((i, j)).map_err(|e| e.to_string())?;
                let allow: Vec<f32> = (0..n_rel)
                    .map(|c| if type_allowed(spec, &mask, &p.facets[i], &p.facets[j], c) { 0.0 } else { f32::NEG_INFINITY })
                    .collect();
                let allow = Tensor::from_vec(allow, n_rel, &device).map_err(|e| e.to_string())?;
                let pred = softmax(&(row + allow).map_err(|e| e.to_string())?, D::Minus1)
                    .and_then(|t| t.argmax(D::Minus1))
                    .and_then(|t| t.to_scalar::<u32>())
                    .map_err(|e| e.to_string())? as usize;
                let want = p.targets[i][j];
                n += 1;
                if pred == want {
                    ok += 1;
                }
                if want != 0 {
                    n_pos += 1;
                    if pred == want {
                        ok_pos += 1;
                    }
                }
            }
        }
    }

    let report = RelReport {
        examples: prepared.len(),
        pairs: total_pairs,
        classes: n_rel,
        first_loss: round4(first_loss),
        last_loss: round4(last_loss),
        train_acc: round4(ok as f64 / n.max(1) as f64),
        train_acc_positive: round4(ok_pos as f64 / n_pos.max(1) as f64),
        dev_examples: dev.len(),
        dev_acc: round4(dev_acc),
        dev_acc_positive: round4(dev_acc_pos),
    };
    Ok((varmap, report))
}

fn round4(v: f64) -> f64 {
    (v * 10000.0).round() / 10000.0
}

/// Persist Head C beside the tagger checkpoint.
pub fn save(varmap: &VarMap, spec: &VocabularySpace, out_dir: &std::path::Path) -> Result<(), String> {
    std::fs::create_dir_all(out_dir).map_err(|e| e.to_string())?;
    varmap.save(out_dir.join("relations.safetensors")).map_err(|e| format!("save head C: {e}"))?;
    let meta = serde_json::json!({
        "classes": crate::tagger_data::head_c_labels(spec),
        "relations": spec.relation_facets,
    });
    std::fs::write(out_dir.join("relations.json"), serde_json::to_vec_pretty(&meta).map_err(|e| e.to_string())?)
        .map_err(|e| e.to_string())?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tagger_data::{Case, RelationLabel};
    use crate::vocabulary::{EntityFacet, RelationFacet};

    fn spec() -> VocabularySpace {
        VocabularySpace {
            version: 1,
            corpus: "t".into(),
            entity_facets: vec![
                EntityFacet { name: "org".into(), parent: None, description: "companies".into(), examples: vec![], structural: false },
                EntityFacet { name: "system".into(), parent: None, description: "platforms".into(), examples: vec![], structural: false },
            ],
            relation_facets: vec![RelationFacet { name: "develops".into(), head: "org".into(), tail: "system".into() }],
            gazetteer: vec![],
            metrics: None,
        }
    }

    #[test]
    fn type_mask_makes_reversed_relations_unrepresentable() {
        let s = spec();
        let m = pair_mask(&s);
        // `none` always allowed
        assert!(type_allowed(&s, &m, "system", "org", 0));
        // declared direction allowed
        assert!(type_allowed(&s, &m, "org", "system", 1));
        // reversed direction is NOT scorable — the guarded fragment inside the model
        assert!(!type_allowed(&s, &m, "system", "org", 1));
        // undeclared facet pairing is not scorable either
        assert!(!type_allowed(&s, &m, "org", "org", 1));
    }

    #[test]
    fn pair_targets_are_directional() {
        let s = spec();
        let ex = TaggerExample {
            text: "Boeing develops the MQ-28.".into(),
            spans: vec![
                LabeledSpan { start: 0, end: 6, facet: "org".into(), surface: "Boeing".into(), negated: false, hedged: false },
                LabeledSpan { start: 20, end: 25, facet: "system".into(), surface: "MQ-28".into(), negated: false, hedged: false },
            ],
            relations: vec![RelationLabel { head: 0, tail: 1, name: "develops".into() }],
            case: Case::Normal,
        };
        let t = pair_targets(&s, &ex);
        assert_eq!(t[0][1], 1, "org→system carries the relation");
        assert_eq!(t[1][0], 0, "system→org is `none`");
    }

    #[test]
    fn biaffine_shapes_and_pooling() {
        let device = Device::Cpu;
        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        let (hp, n_rel, s) = (8usize, 3usize, 4usize);
        let head = BiaffineHead::new(vb, hp, n_rel).unwrap();
        let spans = Tensor::rand(0f32, 1f32, (s, hp), &device).unwrap();
        let logits = head.forward(&spans).unwrap();
        assert_eq!(logits.dims(), &[s, s, n_rel]);

        // pooling gives 2H and respects span/token overlap
        let hidden = Tensor::rand(0f32, 1f32, (1, 5, 4), &device).unwrap();
        let offsets = [(0, 0), (0, 6), (7, 15), (16, 21), (0, 0)];
        let sp = LabeledSpan { start: 0, end: 6, facet: "org".into(), surface: "Boeing".into(), negated: false, hedged: false };
        let pooled = pool_span(&hidden, &offsets, &sp).unwrap();
        assert_eq!(pooled.dims(), &[8]); // 2 * H
        // a span outside every token still yields the right width (zeros)
        let far = LabeledSpan { start: 900, end: 905, facet: "org".into(), surface: "x".into(), negated: false, hedged: false };
        assert_eq!(pool_span(&hidden, &offsets, &far).unwrap().dims(), &[8]);
    }
}