clawmates/crates/cm-decide/src/nli.rs

//! A local decision model: a DeBERTa-v3 MNLI cross-encoder, in-process.
//!
//! A Noul IS an entailment probability — P(the hypothesis follows from the
//! state) — and natural-language inference is the oldest working form of
//! zero-shot classification. A Choice is one hypothesis per option with a
//! softmax over their entailment logits; a Score is the same over ordered
//! levels, then the probability-weighted position. Nothing here reasons or
//! generates: one forward pass per (state, hypothesis) pair, the answer read
//! from the three-way head.
//!
//! What this does NOT have is the vendor's calibration training. Its
//! probabilities are as calibrated as MNLI made them, which on our domain is
//! an open question until `decide-eval` answers it; temperature scaling on
//! our own labelled cases is the standard fix and the `temperature` knob is
//! where it goes.
//!
//! Model files come from the Hugging Face hub on first use
//! (`CLAWMATES_NLI_MODEL`, default `cross-encoder/nli-deberta-v3-base`;
//! `-xsmall` is 4× faster and worse) or from `CLAWMATES_NLI_MODEL_DIR` for a
//! host with no egress. CPU by default; `metal`/`cuda` features select a GPU.

use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::Mutex;
use std::time::Instant;

use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::debertav2::{Config, DebertaV2SeqClassificationModel, Id2Label};
use tokenizers::{EncodeInput, InputSequence, PaddingParams, Tokenizer, TruncationParams};

use crate::{confidence_of, softmax, Answer, DecideError, Decider, Decision, Question};

pub const DEFAULT_MODEL: &str = "cross-encoder/nli-deberta-v3-base";

/// Token budget per pair. DeBERTa-v3 was trained at 512; the STATE is what
/// gets cut when a pair is too long (`OnlyFirst`), so the hypothesis — the
/// question — always survives whole.
const MAX_TOKENS: usize = 512;

pub struct Nli {
    model: Mutex<DebertaV2SeqClassificationModel>,
    tokenizer: Tokenizer,
    device: Device,
    /// Index of each MNLI label in the head, read from `config.json`.
    entail: usize,
    contra: usize,
    /// Divides the logits before the softmax. 1.0 = the model's own
    /// calibration; fitted on labelled cases by the eval when it disagrees.
    temperature: f64,
    name: String,
}

pub struct NliConfig {
    pub model: String,
    pub dir: Option<PathBuf>,
    pub temperature: f64,
}

impl NliConfig {
    pub fn from_env() -> Self {
        Self {
            model: std::env::var("CLAWMATES_NLI_MODEL")
                .ok()
                .filter(|m| !m.trim().is_empty())
                .unwrap_or_else(|| DEFAULT_MODEL.to_string()),
            dir: std::env::var("CLAWMATES_NLI_MODEL_DIR").ok().map(PathBuf::from),
            temperature: std::env::var("CLAWMATES_NLI_TEMPERATURE")
                .ok()
                .and_then(|t| t.parse().ok())
                .unwrap_or(1.0),
        }
    }
}

impl Nli {
    pub fn load(cfg: NliConfig) -> Result<Self, DecideError> {
        let (config_path, tokenizer_path, weights_path) = match &cfg.dir {
            Some(dir) => (
                dir.join("config.json"),
                dir.join("tokenizer.json"),
                dir.join("model.safetensors"),
            ),
            None => {
                let api = hf_hub::api::sync::Api::new()
                    .map_err(|e| DecideError::Config(format!("hf hub: {e}")))?;
                let repo = api.model(cfg.model.clone());
                let get = |f: &str| {
                    repo.get(f)
                        .map_err(|e| DecideError::Config(format!("fetch {f} for {}: {e}", cfg.model)))
                };
                (get("config.json")?, get("tokenizer.json")?, get("model.safetensors")?)
            }
        };
        let config_text = std::fs::read_to_string(&config_path)
            .map_err(|e| DecideError::Config(format!("{}: {e}", config_path.display())))?;
        let config: Config = serde_json::from_str(&config_text)
            .map_err(|e| DecideError::Config(format!("config.json: {e}")))?;
        // The head's label order is the model's, not ours: read it.
        let id2label: Id2Label = serde_json::from_str::<serde_json::Value>(&config_text)
            .ok()
            .and_then(|v| v.get("id2label").cloned())
            .and_then(|v| {
                v.as_object().map(|m| {
                    m.iter()
                        .filter_map(|(k, v)| {
                            Some((k.parse::<u32>().ok()?, v.as_str()?.to_lowercase()))
                        })
                        .collect()
                })
            })
            .ok_or_else(|| DecideError::Config("config.json has no id2label".into()))?;
        let find = |name: &str| {
            id2label
                .iter()
                .find(|(_, v)| v.as_str() == name)
                .map(|(k, _)| *k as usize)
                .ok_or_else(|| DecideError::Config(format!("head has no {name:?} label")))
        };
        let entail = find("entailment")?;
        // MNLI heads have three labels; the zero-shot-tuned checkpoints
        // collapse to entailment / not_entailment. Either "no" label works.
        let contra = find("contradiction").or_else(|_| find("not_entailment"))?;

        let device = pick_device();
        // Buffered, not mmap'd: the workspace denies `unsafe`, and the weights
        // (740 MB for -base) are read once into memory and kept for the life
        // of the process anyway.
        let bytes = std::fs::read(&weights_path)
            .map_err(|e| DecideError::Config(format!("{}: {e}", weights_path.display())))?;
        let vb = VarBuilder::from_buffered_safetensors(bytes, DType::F32, &device)
            .map_err(|e| DecideError::Config(format!("weights: {e}")))?;
        // The backbone's tensors are `deberta.*`; the pooler and classifier
        // read from the root. candle's own example does exactly this.
        let vb = vb.set_prefix("deberta");
        let model = DebertaV2SeqClassificationModel::load(vb, &config, Some(id2label))
            .map_err(|e| DecideError::Config(format!("model: {e}")))?;

        let mut tokenizer = Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| DecideError::Config(format!("tokenizer: {e}")))?;
        tokenizer.with_padding(Some(PaddingParams::default()));
        tokenizer
            .with_truncation(Some(TruncationParams {
                max_length: MAX_TOKENS,
                strategy: tokenizers::TruncationStrategy::OnlyFirst,
                ..Default::default()
            }))
            .map_err(|e| DecideError::Config(format!("truncation: {e}")))?;

        Ok(Self {
            model: Mutex::new(model),
            tokenizer,
            device,
            entail,
            contra,
            temperature: cfg.temperature.max(1e-3),
            name: format!("nli:{}", cfg.model.rsplit('/').next().unwrap_or(&cfg.model)),
        })
    }

    /// Entailment and contradiction logits for every (state, hypothesis)
    /// pair, one batch, one forward pass.
    fn logits(&self, state: &str, hypotheses: &[String]) -> Result<Vec<(f64, f64)>, DecideError> {
        if hypotheses.is_empty() {
            return Ok(Vec::new());
        }
        let inputs: Vec<EncodeInput> = hypotheses
            .iter()
            .map(|h| {
                EncodeInput::Dual(
                    InputSequence::Raw(state.into()),
                    InputSequence::Raw(h.as_str().into()),
                )
            })
            .collect();
        let encodings = self
            .tokenizer
            .encode_batch(inputs, true)
            .map_err(|e| DecideError::Model(format!("tokenize: {e}")))?;
        let to_tensor = |rows: Vec<&[u32]>| -> Result<Tensor, DecideError> {
            let stacked: Vec<Tensor> = rows
                .iter()
                .map(|r| Tensor::new(*r, &self.device))
                .collect::<Result<_, _>>()
                .map_err(|e| DecideError::Model(e.to_string()))?;
            Tensor::stack(&stacked, 0).map_err(|e| DecideError::Model(e.to_string()))
        };
        let ids = to_tensor(encodings.iter().map(|e| e.get_ids()).collect())?;
        let mask = to_tensor(encodings.iter().map(|e| e.get_attention_mask()).collect())?;
        let types = to_tensor(encodings.iter().map(|e| e.get_type_ids()).collect())?;
        let out = {
            let model = self.model.lock().map_err(|_| DecideError::Model("model lock".into()))?;
            model
                .forward(&ids, Some(types), Some(mask))
                .map_err(|e| DecideError::Model(format!("forward: {e}")))?
        };
        let rows: Vec<Vec<f32>> = out
            .to_dtype(DType::F32)
            .and_then(|t| t.to_vec2())
            .map_err(|e| DecideError::Model(e.to_string()))?;
        Ok(rows
            .iter()
            .map(|r| {
                (
                    r[self.entail] as f64 / self.temperature,
                    r[self.contra] as f64 / self.temperature,
                )
            })
            .collect())
    }

    /// P(yes) for one hypothesis: entailment against contradiction, the
    /// neutral class left out — the standard zero-shot NLI reading.
    fn p_yes(entail: f64, contra: f64) -> f64 {
        softmax(&[entail, contra])[0]
    }
}

fn pick_device() -> Device {
    #[cfg(feature = "cuda")]
    if let Ok(d) = Device::new_cuda(0) {
        return d;
    }
    #[cfg(feature = "metal")]
    if let Ok(d) = Device::new_metal(0) {
        return d;
    }
    Device::Cpu
}

/// The hypothesis a question becomes. Kept in one place because the wording
/// is the whole "prompt" this backend has, and the eval is what tunes it.
pub fn hypotheses(question: &Question) -> Vec<String> {
    match question {
        Question::Noul { instructions, criteria } => match criteria {
            Some(c) => vec![
                format!("{instructions} {}", c.yes),
                format!("{instructions} {}", c.no),
            ],
            None => vec![instructions.clone()],
        },
        Question::Choice { instructions, criteria } => criteria
            .iter()
            .map(|(name, desc)| match desc {
                Some(d) => format!("{instructions} {name}: {d}"),
                None => format!("{instructions} {name}"),
            })
            .collect(),
        Question::Score { instructions, criteria } => criteria
            .iter()
            .map(|level| format!("{instructions} {level}"))
            .collect(),
    }
}

#[async_trait::async_trait]
impl Decider for Nli {
    fn name(&self) -> &str {
        &self.name
    }

    async fn decide(
        &self,
        state: &str,
        questions: &BTreeMap<String, Question>,
    ) -> Result<Decision, DecideError> {
        let started = Instant::now();
        // Every hypothesis of every question in one batch; then read each
        // question's slice back out.
        let mut all: Vec<String> = Vec::new();
        let mut spans: Vec<(String, usize, usize)> = Vec::new();
        for (id, q) in questions {
            let hs = hypotheses(q);
            let start = all.len();
            all.extend(hs);
            spans.push((id.clone(), start, all.len()));
        }
        let logits = self.logits(state, &all)?;
        let mut answers = BTreeMap::new();
        for (id, start, end) in spans {
            let slice = &logits[start..end];
            let q = &questions[&id];
            let answer = match q {
                Question::Noul { criteria, .. } => {
                    let p = if criteria.is_some() {
                        // yes-description vs no-description: which does the
                        // state entail more?
                        softmax(&[slice[0].0, slice[1].0])[0]
                    } else {
                        Self::p_yes(slice[0].0, slice[0].1)
                    };
                    Answer::Noul { noul: p }
                }
                Question::Choice { criteria, .. } => {
                    let probs = softmax(&slice.iter().map(|(e, _)| *e).collect::<Vec<_>>());
                    let names: Vec<&String> = criteria.keys().collect();
                    let (best, _) = probs
                        .iter()
                        .enumerate()
                        .fold((0, f64::NEG_INFINITY), |acc, (i, p)| if *p > acc.1 { (i, *p) } else { acc });
                    Answer::Choice {
                        choice: names[best].clone(),
                        confidence: confidence_of(&probs),
                        probabilities: names.iter().map(|n| (*n).clone()).zip(probs).collect(),
                    }
                }
                Question::Score { .. } => {
                    let probs = softmax(&slice.iter().map(|(e, _)| *e).collect::<Vec<_>>());
                    let score = probs.iter().enumerate().map(|(i, p)| i as f64 * p).sum();
                    Answer::Score {
                        score,
                        confidence: confidence_of(&probs),
                        probabilities: probs,
                    }
                }
            };
            answers.insert(id, answer);
        }
        Ok(Decision {
            model: self.name.clone(),
            answers,
            usage: None,
            latency: started.elapsed(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn every_question_becomes_the_expected_hypotheses() {
        let n = Question::noul("The task needs web search.");
        assert_eq!(hypotheses(&n), vec!["The task needs web search."]);
        let c = Question::choice("This work is", [("research", Some("reading")), ("coding", None)]);
        let h = hypotheses(&c);
        assert_eq!(h, vec!["This work is coding", "This work is research: reading"]);
        let s = Question::score("Severity:", ["cosmetic", "blocking"]);
        assert_eq!(hypotheses(&s).len(), 2);
    }

    #[test]
    fn p_yes_is_entailment_against_contradiction() {
        assert!(Nli::p_yes(3.0, -3.0) > 0.99);
        assert!(Nli::p_yes(-3.0, 3.0) < 0.01);
        assert!((Nli::p_yes(0.0, 0.0) - 0.5).abs() < 1e-9);
    }
}