bra0-kg 0.2.3

bra0 Knowledge Graph core — RDF transformations (sophia) + KgStore trait (NextGraph-First)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223

//! Zero-shot NER via candle BERT + SKOS vocabulary
//!
//! Capability: NS-8 (Entity Recognition)
//! Tier 2 of the symbolic-first cascade (NS-5)
//!
//! Pipeline:
//!   1. Input: raw text + SKOS concept labels (from KG vocabulary)
//!   2. Encode text and each candidate label via BERT sentence embedding
//!   3. Cosine similarity ranks candidates
//!   4. Output: NerResult[] — matched concepts with confidence scores
//!
//! Model: sentence-transformers/all-MiniLM-L6-v2 (22MB, 6 layers, 384 dims)
//! WASM: 4.1 MB binary, 97MB runtime, 1546ms inference (proven SPIKE-02)

use candle_core::{Device, Tensor, DType};
use candle_nn::VarBuilder;
use candle_transformers::models::bert::{BertModel, Config};
use tokenizers::Tokenizer;
use serde::{Deserialize, Serialize};

/// A candidate entity from the domain SKOS vocabulary.
#[derive(Debug, Clone, Deserialize)]
pub struct Candidate {
    pub label: String,
    pub iri: String,
}

/// A NER match result.
#[derive(Debug, Clone, Serialize)]
pub struct NerResult {
    pub label: String,
    pub iri: String,
    pub similarity: f32,
}

/// BERT-based zero-shot NER engine.
///
/// Loads a sentence-transformer model and computes cosine similarity
/// between input text and SKOS candidate labels. No fine-tuning needed —
/// the SKOS vocabulary IS the label set.
pub struct NerEngine {
    model: BertModel,
    tokenizer: Tokenizer,
}

impl NerEngine {
    /// Load model from in-memory buffers (WASM-compatible — no filesystem).
    ///
    /// `weights`: SafeTensors bytes (fetched by JS in browser, or read from disk native)
    /// `tokenizer_json`: tokenizer.json content
    /// `config_json`: config.json content
    pub fn from_bytes(
        weights: &[u8],
        tokenizer_json: &str,
        config_json: &str,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let device = Device::Cpu;

        let tokenizer = Tokenizer::from_bytes(tokenizer_json.as_bytes())
            .map_err(|e| format!("tokenizer: {}", e))?;

        let config: Config = serde_json::from_str(config_json)?;

        let vb = VarBuilder::from_buffered_safetensors(
            weights.to_vec(),
            DType::F32,
            &device,
        )?;

        let model = BertModel::load(vb, &config)?;
        Ok(Self { model, tokenizer })
    }

    /// Load model from HuggingFace Hub (native only — not WASM).
    #[cfg(not(target_arch = "wasm32"))]
    pub fn from_hf(model_id: &str) -> Result<Self, Box<dyn std::error::Error>> {
        let api = hf_hub::api::sync::Api::new()?;
        let repo = api.model(model_id.to_string());

        let tokenizer_path = repo.get("tokenizer.json")?;
        let tokenizer_json = std::fs::read_to_string(tokenizer_path)?;

        let config_path = repo.get("config.json")?;
        let config_json = std::fs::read_to_string(config_path)?;

        let weights_path = repo.get("model.safetensors")?;
        let weights = std::fs::read(weights_path)?;

        Self::from_bytes(&weights, &tokenizer_json, &config_json)
    }

    /// Encode a text into a sentence embedding (mean pooling).
    pub fn encode(&self, text: &str) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
        let encoding = self.tokenizer.encode(text, true)
            .map_err(|e| format!("tokenize: {}", e))?;
        let ids = encoding.get_ids();
        let token_ids = Tensor::new(ids, &Device::Cpu)?.unsqueeze(0)?;
        let type_ids = token_ids.zeros_like()?;

        let embeddings = self.model.forward(&token_ids, &type_ids, None)?;

        let (_, seq_len, _) = embeddings.dims3()?;
        let mean = (embeddings.sum(1)? / (seq_len as f64))?.squeeze(0)?;
        Ok(mean.to_vec1::<f32>()?)
    }

    /// Zero-shot NER: rank SKOS candidates by cosine similarity with text.
    ///
    /// Returns candidates sorted by similarity descending.
    /// Candidates with similarity > `threshold` are considered matches.
    pub fn zero_shot_ner(
        &self,
        text: &str,
        candidates: &[Candidate],
        threshold: f32,
    ) -> Result<Vec<NerResult>, Box<dyn std::error::Error>> {
        let text_emb = self.encode(text)?;

        let mut results: Vec<NerResult> = candidates.iter()
            .map(|c| {
                let label_emb = self.encode(&c.label)?;
                let sim = cosine_similarity(&text_emb, &label_emb);
                Ok(NerResult {
                    label: c.label.clone(),
                    iri: c.iri.clone(),
                    similarity: sim,
                })
            })
            .collect::<Result<Vec<_>, Box<dyn std::error::Error>>>()?;

        results.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap());

        if threshold > 0.0 {
            results.retain(|r| r.similarity >= threshold);
        }

        Ok(results)
    }
}

fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm_a == 0.0 || norm_b == 0.0 { return 0.0; }
    dot / (norm_a * norm_b)
}

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

    /// NS-8 acceptance: load MiniLM, run NER on medical text, top match is relevant.
    #[test]
    #[cfg(not(target_arch = "wasm32"))]
    fn test_zero_shot_ner_medical() {
        // Skip if no network / model not cached
        let engine = match NerEngine::from_hf("sentence-transformers/all-MiniLM-L6-v2") {
            Ok(e) => e,
            Err(e) => { eprintln!("SKIP NER test: {}", e); return; }
        };

        let text = "The patient presents with acute myocardial infarction and requires immediate percutaneous coronary intervention.";

        let candidates = vec![
            Candidate { label: "myocardial infarction".into(), iri: "http://snomed.info/id/22298006".into() },
            Candidate { label: "percutaneous coronary intervention".into(), iri: "http://snomed.info/id/415070008".into() },
            Candidate { label: "diabetes mellitus".into(), iri: "http://snomed.info/id/73211009".into() },
            Candidate { label: "fracture".into(), iri: "http://snomed.info/id/125605004".into() },
        ];

        let results = engine.zero_shot_ner(text, &candidates, 0.0).unwrap();

        eprintln!("NER results:");
        for r in &results {
            eprintln!("  {:.4}  {} ({})", r.similarity, r.label, r.iri);
        }

        // Top 2 should be the relevant medical terms
        assert!(results[0].label.contains("coronary") || results[0].label.contains("myocardial"),
            "Top match should be a relevant medical term, got: {}", results[0].label);
        assert!(results[0].similarity > 0.85,
            "Top similarity should be > 0.85, got: {}", results[0].similarity);

        // Threshold filtering
        let filtered = engine.zero_shot_ner(text, &candidates, 0.9).unwrap();
        assert!(filtered.len() <= 2, "Only top matches should survive 0.9 threshold");
    }

    /// Latency gate: < 2s in release, validated by SPIKE-02 Puppeteer (1546ms WASM).
    /// IGNORED in debug (candle is ~100x slower unoptimized).
    /// Run with: cargo test --release --features ner -- test_ner_latency
    #[test]
    #[ignore]
    #[cfg(not(target_arch = "wasm32"))]
    fn test_ner_latency() {
        let engine = match NerEngine::from_hf("sentence-transformers/all-MiniLM-L6-v2") {
            Ok(e) => e,
            Err(e) => { eprintln!("SKIP latency test: {}", e); return; }
        };

        // Warm up — first encode is slow (model initialization)
        let _ = engine.encode("warmup");

        let text = "Chronic obstructive pulmonary disease with acute exacerbation requiring hospitalization.";
        let candidates = vec![
            Candidate { label: "COPD".into(), iri: "http://snomed.info/id/13645005".into() },
            Candidate { label: "asthma".into(), iri: "http://snomed.info/id/195967001".into() },
            Candidate { label: "pneumonia".into(), iri: "http://snomed.info/id/233604007".into() },
        ];

        let start = std::time::Instant::now();
        let results = engine.zero_shot_ner(text, &candidates, 0.0).unwrap();
        let elapsed = start.elapsed();

        eprintln!("NER inference latency: {:?} for {} candidates (after warmup)", elapsed, candidates.len());
        // 2s gate for release/WASM. Debug is ~5x slower but should still be < 30s.
        let gate_ms: u128 = if cfg!(debug_assertions) { 30000 } else { 2000 };
        assert!(elapsed.as_millis() < gate_ms,
            "NER inference should complete in < {}ms, took {:?}", gate_ms, elapsed);
        assert!(!results.is_empty());
    }
}