//! **Step 1: the tagger's finetuning dataset** — synthetic, spec-driven, and validated by the engine.
//!
//! ## Head design (locked here, because the dataset's shape *is* the head design)
//!
//! The discovered [`VocabularySpace`] drives every head's output space. This is the load-bearing
//! property: a model whose label set is *derived from* the spec cannot emit an out-of-vocabulary facet or
//! a type-invalid relation, so its output is linter-clean by construction (paper §2).
//!
//! ```text
//! shared encoder (mmBERT-small H=384 / bert-tiny H=128)
//! │ last_hidden_state [T,H]
//! ┌──────────────────────────┼──────────────────────────┐
//! Head A: BIO span typing Head B: epistemic span pooling
//! [T,H] → [T, 2K+1] [T,H] → [T,4] (start⊕end⊕mean per span)
//! K = spec entity facets │
//! + REL/QTY/GEO/TIME asserted / hedged / Head C: biaffine relation
//! → dims 1, 3, 4 negated / negated+hedged [S,H]×[S,H] → [S,S,R+1]
//! → dim 5, and infon R = spec relation facets
//! polarity i ∈ {±1, ±0.5} → dim 2 (polarity per argument)
//! ```
//!
//! * **Head A** emits the *typed* span so dimension 1 is real (`org/…` not `ent/…`). Its label set is
//! [`head_a_labels`].
//! * **Head B** emits the epistemic reading, which is simultaneously dimension 5 (`state/negated`) and
//! the Dempster-Shafer polarity `i` fed to `InfonIndex::add_infon_polar` (see
//! [`crate::dimensions::belief_level`]).
//! * **Head C** binds arguments. Polarity is a property of the *argument side*, not of the predicate
//! token, so a predicted pair `(h,t)` for relation `r` emits `rel/r/+` on `h` and `rel/r/-` on `t`.
//! Crucially its scores are **masked by the spec's declared `head`/`tail` facets**
//! ([`pair_mask`]): only type-valid pairs are scorable. That shrinks the `O(S²)` pair space, enforces
//! the guarded fragment at the model level, and makes an invalid relation unrepresentable.
//!
//! ## Dataset contract
//!
//! One JSON object per line: `text`, char-offset `spans` (each with a spec facet), per-span `epistemic`
//! flags, and `relations` as span-index pairs. Generation is LLM-driven but every example is
//! **mechanically validated** ([`validate`]) before it is kept: surfaces must align to exact offsets,
//! facets must be declared, relations must satisfy the spec's head/tail types, and spans must not
//! overlap. Coverage of the hard cases (coref, hedged, negated, adversarial no-relation) is planned
//! explicitly by [`GenPlan`], not left to chance.
use crate::vocabulary::VocabularySpace;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
/// Non-entity span kinds every tagger carries regardless of the corpus (dimensions 3-4 plus the relation
/// predicate itself).
pub const STRUCTURAL_KINDS: &[&str] = &["REL", "QTY", "GEO", "TIME", "STATE"];
/// A labelled span: byte offsets into `text` plus its spec facet.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct LabeledSpan {
pub start: usize,
pub end: usize,
/// a spec entity-facet name, or one of [`STRUCTURAL_KINDS`]
pub facet: String,
pub surface: String,
/// Head B target for this span: was the assertion negated / hedged?
#[serde(default)]
pub negated: bool,
#[serde(default)]
pub hedged: bool,
}
/// A bound relation: span indices + a spec relation name. Head C's target.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RelationLabel {
pub head: usize,
pub tail: usize,
pub name: String,
}
/// Which hard case an example exercises — tracked so the generator can guarantee coverage instead of
/// hoping for it.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Case {
/// a plain, directly-stated relation
Normal,
/// an argument referred to by pronoun or a definite description ("it", "the vehicle")
Coref,
/// hedged assertion ("may", "reportedly") → weak belief
Hedged,
/// explicit negation → negative belief
Negated,
/// entities co-occur but assert NO relation — the hard negative Head C needs
Adversarial,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaggerExample {
pub text: String,
pub spans: Vec<LabeledSpan>,
#[serde(default)]
pub relations: Vec<RelationLabel>,
pub case: Case,
}
// ── head output spaces, derived from the spec ───────────────────────────────────────────────────
/// Head A's BIO label set: `O` + `B-`/`I-` for every spec entity facet and structural kind. Order is
/// deterministic (spec order, then structural) so a checkpoint's label indices stay stable.
pub fn head_a_labels(spec: &VocabularySpace) -> Vec<String> {
let mut out = vec!["O".to_string()];
for f in spec.taggable_facets().iter().map(|f| f.name.to_uppercase()).chain(STRUCTURAL_KINDS.iter().map(|s| s.to_string())) {
out.push(format!("B-{f}"));
out.push(format!("I-{f}"));
}
out
}
/// Head B's label set — the four epistemic readings, in polarity order (`+1, +0.5, -0.5, -1`).
pub fn head_b_labels() -> [&'static str; 4] {
["asserted", "hedged", "negated_hedged", "negated"]
}
/// Head C's label set: `none` plus every declared relation.
pub fn head_c_labels(spec: &VocabularySpace) -> Vec<String> {
let mut out = vec!["none".to_string()];
out.extend(spec.relation_facets.iter().map(|r| r.name.clone()));
out
}
/// Type-valid `(head_facet, tail_facet, relation)` triples — the mask applied to Head C's pair scores.
/// Any pair outside this set is not scorable, which is what makes a type-invalid relation
/// unrepresentable rather than merely unlikely.
pub fn pair_mask(spec: &VocabularySpace) -> HashSet<(String, String, String)> {
spec.relation_facets.iter().map(|r| (r.head.clone(), r.tail.clone(), r.name.clone())).collect()
}
// ── alignment: surfaces → exact byte spans (the reference `bio_for`/`paint` step) ────────────────
/// Find the first occurrence of `surface` in `text` not overlapping `taken`, case-insensitively but
/// returning offsets into the original text. Returns `None` if the surface isn't present — the signal
/// that a generated example must be discarded rather than silently mislabelled.
///
/// Matches are **word-boundary anchored**: without this, a short anaphor like `it` aligns to the `it`
/// inside `submitted`, silently mislabelling a training example. Boundaries are required on whichever
/// side the surface itself is alphanumeric, so punctuated surfaces (`F-35`, `MQ-28,`) still match.
pub fn align(text: &str, surface: &str, taken: &[(usize, usize)]) -> Option<(usize, usize)> {
// Case-insensitive match walked over the ORIGINAL text. Searching a `to_lowercase()` copy and reusing
// its byte offsets is wrong: lowercasing can change a character's byte length, so offsets drift in any
// passage containing non-ASCII — which silently mislabels spans (observed on ~11% of real passages).
let needle: Vec<char> = surface.trim().chars().flat_map(|c| c.to_lowercase()).collect();
if needle.is_empty() {
return None;
}
let chars: Vec<(usize, char)> = text.char_indices().collect();
let needle_first_alnum = needle.first().map(|c| c.is_alphanumeric()).unwrap_or(false);
let needle_last_alnum = needle.last().map(|c| c.is_alphanumeric()).unwrap_or(false);
for si in 0..chars.len() {
let mut ni = 0usize;
let mut ci = si;
let mut matched = true;
while ni < needle.len() && ci < chars.len() {
let mut consumed_all = true;
for lc in chars[ci].1.to_lowercase() {
if ni < needle.len() && needle[ni] == lc {
ni += 1;
} else {
consumed_all = false;
break;
}
}
if !consumed_all {
matched = false;
break;
}
ci += 1;
}
if !matched || ni != needle.len() {
continue;
}
let start = chars[si].0;
let end = if ci < chars.len() { chars[ci].0 } else { text.len() };
// word boundaries, required only on sides where the surface itself is alphanumeric
let before_ok = !needle_first_alnum || si == 0 || !chars[si - 1].1.is_alphanumeric();
let after_ok = !needle_last_alnum || ci >= chars.len() || !chars[ci].1.is_alphanumeric();
if before_ok && after_ok && !taken.iter().any(|(ts, te)| start < *te && *ts < end) {
return Some((start, end));
}
}
None
}
/// Project char-span labels onto tokenizer offsets as BIO indices for Head A. `offsets` are
/// `(start, end)` byte ranges per token (as the `tokenizers` crate reports); `-100` marks ignored
/// positions (specials / zero-width), matching the training convention.
pub fn to_bio(spec: &VocabularySpace, spans: &[LabeledSpan], offsets: &[(usize, usize)]) -> Vec<i64> {
let labels = head_a_labels(spec);
let idx = |l: &str| labels.iter().position(|x| x == l).map(|i| i as i64).unwrap_or(0);
let mut out = vec![0i64; offsets.len()];
for (ti, (ts, te)) in offsets.iter().enumerate() {
if te <= ts {
out[ti] = -100; // special / empty token
continue;
}
if let Some(sp) = spans.iter().find(|s| *ts < s.end && s.start < *te) {
let kind = sp.facet.to_uppercase();
let first = *ts <= sp.start;
out[ti] = idx(&format!("{}-{}", if first { "B" } else { "I" }, kind));
}
}
out
}
/// Head B target per token: the epistemic class of the span covering it (`asserted` elsewhere).
pub fn to_epistemic(spans: &[LabeledSpan], offsets: &[(usize, usize)]) -> Vec<i64> {
let class = |s: &LabeledSpan| match (s.negated, s.hedged) {
(false, false) => 0, // asserted
(false, true) => 1, // hedged
(true, true) => 2, // negated_hedged
(true, false) => 3, // negated
};
offsets
.iter()
.map(|(ts, te)| {
if te <= ts {
return -100;
}
spans.iter().find(|s| *ts < s.end && s.start < *te).map(class).unwrap_or(0)
})
.collect()
}
// ── validation: the engine is the oracle ────────────────────────────────────────────────────────
#[derive(Debug, Clone, PartialEq)]
pub enum Reject {
NoSpans,
SurfaceNotFound(String),
UndeclaredFacet(String),
OverlappingSpans,
BadRelationIndex,
UndeclaredRelation(String),
/// the relation's arguments don't match the spec's declared head/tail facets
TypeMismatch { name: String, got: (String, String), want: (String, String) },
}
/// Mechanically check an example against the spec. Everything here is a *hard* reject: a mislabelled
/// training example is worse than a missing one.
pub fn validate(spec: &VocabularySpace, ex: &TaggerExample) -> Result<(), Reject> {
if ex.spans.is_empty() {
return Err(Reject::NoSpans);
}
for s in &ex.spans {
let declared = spec.has_entity_facet(&s.facet) || STRUCTURAL_KINDS.contains(&s.facet.to_uppercase().as_str());
if !declared {
return Err(Reject::UndeclaredFacet(s.facet.clone()));
}
if s.end > ex.text.len() || s.start >= s.end {
return Err(Reject::SurfaceNotFound(s.surface.clone()));
}
if ex.text[s.start..s.end].to_lowercase() != s.surface.trim().to_lowercase() {
return Err(Reject::SurfaceNotFound(s.surface.clone()));
}
}
// no overlaps (Head A is single-label BIO)
let mut sorted: Vec<&LabeledSpan> = ex.spans.iter().collect();
sorted.sort_by_key(|s| s.start);
if sorted.windows(2).any(|w| w[0].end > w[1].start) {
return Err(Reject::OverlappingSpans);
}
for r in &ex.relations {
let (Some(h), Some(t)) = (ex.spans.get(r.head), ex.spans.get(r.tail)) else {
return Err(Reject::BadRelationIndex);
};
// A relation NAME may be declared with several type signatures — real ontologies do this
// (`has_type: species → type` and `has_type: move → type`). Accept when ANY declaration matches;
// looking at only the first would reject legitimate overloads.
let decls: Vec<&crate::vocabulary::RelationFacet> = spec.relation_facets.iter().filter(|d| d.name == r.name).collect();
if decls.is_empty() {
return Err(Reject::UndeclaredRelation(r.name.clone()));
}
if !decls.iter().any(|d| d.head == h.facet && d.tail == t.facet) {
let first = decls[0];
return Err(Reject::TypeMismatch {
name: r.name.clone(),
got: (h.facet.clone(), t.facet.clone()),
want: (first.head.clone(), first.tail.clone()),
});
}
}
Ok(())
}
// ── generation plan + prompts ───────────────────────────────────────────────────────────────────
/// How many examples of each hard case to request per relation. Defaults mirror the reference
/// generator's ~1-in-8 negative rate and guarantee the coref/hedged/negated coverage the heads need.
#[derive(Debug, Clone, Copy)]
pub struct GenPlan {
pub normal: usize,
pub coref: usize,
pub hedged: usize,
pub negated: usize,
pub adversarial: usize,
}
impl Default for GenPlan {
fn default() -> Self {
GenPlan { normal: 4, coref: 2, hedged: 2, negated: 2, adversarial: 2 }
}
}
impl GenPlan {
pub fn total_per_relation(&self) -> usize {
self.normal + self.coref + self.hedged + self.negated + self.adversarial
}
pub fn cases(&self) -> Vec<(Case, usize)> {
vec![
(Case::Normal, self.normal),
(Case::Coref, self.coref),
(Case::Hedged, self.hedged),
(Case::Negated, self.negated),
(Case::Adversarial, self.adversarial),
]
}
}
/// Instruction for one case — spelled out because these distinctions are exactly what the heads learn.
pub fn case_instruction(case: Case) -> &'static str {
match case {
Case::Normal => "State the relation directly and plainly.",
Case::Coref => "Write TWO clauses: name the argument in the first, then refer back to it with a PRONOUN or short anaphor (\"it\", \"they\", \"the aircraft\") in the second, where the relation is asserted. Label the ANAPHOR as the span (not the earlier mention), so the model must resolve the reference.",
Case::Hedged => "Hedge the assertion (\"may\", \"is expected to\", \"reportedly\"). Set hedged=true on the argument spans AND add one span covering the hedge cue itself with facet \"state\" and hedged=true.",
Case::Negated => "Explicitly negate the relation (\"does not\", \"never\", \"was not\"). Set negated=true on the argument spans AND add one span covering the negation cue itself with facet \"state\" and negated=true.",
Case::Adversarial => "Mention both entity types in one sentence but assert NO relation between them (they merely co-occur). Return an empty relations list.",
}
}
/// The structured-output schema the generator model must fill.
pub fn generation_schema() -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": { "examples": { "type": "array", "items": { "type": "object", "properties": {
"text": {"type": "string"},
"spans": {"type": "array", "items": {"type": "object", "properties": {
"surface": {"type": "string"}, "facet": {"type": "string"},
"negated": {"type": "boolean"}, "hedged": {"type": "boolean"}
}, "required": ["surface", "facet"]}},
"relations": {"type": "array", "items": {"type": "object", "properties": {
"head_surface": {"type": "string"}, "tail_surface": {"type": "string"}, "name": {"type": "string"}
}, "required": ["head_surface", "tail_surface", "name"]}}
}, "required": ["text", "spans"] } } },
"required": ["examples"]
})
}
pub const GENERATION_SYSTEM: &str = "You generate labelled training sentences for a span-tagging and relation-extraction model. Every span's 'surface' MUST appear VERBATIM in 'text' (exact substring, same casing where possible) and its 'facet' MUST be one of the facets given. Relations reference spans by their exact surface strings and must respect the declared head/tail facet types. Write natural domain sentences, one relation per sentence unless told otherwise. Keep sentences under 40 words.";
/// Build the user prompt for one (relation, case) request.
pub fn generation_prompt(spec: &VocabularySpace, relation: &str, case: Case, n: usize) -> String {
let facets: Vec<String> = spec.taggable_facets().iter().map(|f| format!("{} — {}", f.name, f.description)).collect();
let decl = spec.relation(relation);
let rel_line = match decl {
Some(r) => format!("relation '{}': head facet '{}' acts on tail facet '{}'", r.name, r.head, r.tail),
None => format!("relation '{relation}'"),
};
format!(
"Entity facets:\n{}\n\nTarget {rel_line}\n\nGenerate {n} examples. {}\n\nCorpus domain: {}",
facets.join("\n"),
case_instruction(case),
spec.corpus
)
}
/// Convert a model's surface-based proposal into offset-aligned examples, discarding anything that
/// fails alignment or validation. Returns `(kept, rejects)` so generation quality is observable.
pub fn examples_from_proposal(spec: &VocabularySpace, v: &serde_json::Value, case: Case) -> (Vec<TaggerExample>, Vec<Reject>) {
let mut kept = Vec::new();
let mut rejects = Vec::new();
let Some(items) = v.get("examples").and_then(|x| x.as_array()) else { return (kept, rejects) };
for it in items {
let Some(text) = it.get("text").and_then(|t| t.as_str()) else { continue };
let text = text.trim().to_string();
let mut spans: Vec<LabeledSpan> = Vec::new();
let mut taken: Vec<(usize, usize)> = Vec::new();
let mut failed: Option<Reject> = None;
for sp in it.get("spans").and_then(|x| x.as_array()).map(|a| a.as_slice()).unwrap_or(&[]) {
let (Some(surface), Some(facet)) = (sp.get("surface").and_then(|s| s.as_str()), sp.get("facet").and_then(|s| s.as_str())) else { continue };
let facet = crate::projector::slug(facet);
match align(&text, surface, &taken) {
Some((s, e)) => {
taken.push((s, e));
spans.push(LabeledSpan {
start: s,
end: e,
facet,
surface: text[s..e].to_string(),
negated: sp.get("negated").and_then(|b| b.as_bool()).unwrap_or(case == Case::Negated),
hedged: sp.get("hedged").and_then(|b| b.as_bool()).unwrap_or(case == Case::Hedged),
});
}
None => failed = Some(Reject::SurfaceNotFound(surface.to_string())),
}
}
if let Some(r) = failed {
rejects.push(r);
continue;
}
// relations reference spans by surface → resolve to indices
let mut relations = Vec::new();
let find = |s: &str| spans.iter().position(|x| x.surface.to_lowercase() == s.trim().to_lowercase());
for r in it.get("relations").and_then(|x| x.as_array()).map(|a| a.as_slice()).unwrap_or(&[]) {
let (Some(hs), Some(ts), Some(name)) =
(r.get("head_surface").and_then(|s| s.as_str()), r.get("tail_surface").and_then(|s| s.as_str()), r.get("name").and_then(|s| s.as_str()))
else {
continue;
};
match (find(hs), find(ts)) {
(Some(h), Some(t)) => relations.push(RelationLabel { head: h, tail: t, name: crate::projector::slug(name) }),
_ => rejects.push(Reject::BadRelationIndex),
}
}
let ex = TaggerExample { text, spans, relations, case };
match validate(spec, &ex) {
Ok(()) => kept.push(ex),
Err(r) => rejects.push(r),
}
}
(kept, rejects)
}
// ── generation driver ───────────────────────────────────────────────────────────────────────────
/// Coverage report for a generation run — kept counts per case plus reject reasons, so dataset quality
/// is observable rather than assumed.
#[derive(Debug, Default, Serialize)]
pub struct GenReport {
pub kept: usize,
pub rejected: usize,
pub per_case: std::collections::BTreeMap<String, usize>,
pub reject_reasons: std::collections::BTreeMap<String, usize>,
}
/// Drive an LLM over every (relation × hard case) in the plan, aligning and validating each batch.
/// Requests are sequential to keep provider pressure predictable; failures degrade a single batch rather
/// than the run.
#[cfg(feature = "agent")]
pub async fn generate(
provider: &dyn crate::agent::provider::LlmProvider,
spec: &VocabularySpace,
plan: GenPlan,
) -> (Vec<TaggerExample>, GenReport) {
use crate::agent::types::{Msg, ToolSpec};
let tools = vec![ToolSpec {
name: "emit_examples".into(),
description: "Emit labelled training sentences.".into(),
schema: generation_schema(),
}];
let mut out: Vec<TaggerExample> = Vec::new();
let mut rep = GenReport::default();
for rel in &spec.relation_facets {
for (case, n) in plan.cases() {
if n == 0 {
continue;
}
let prompt = generation_prompt(spec, &rel.name, case, n);
let turn = match provider.chat(GENERATION_SYSTEM, &[Msg::user_text(prompt)], &tools).await {
Ok(t) => t,
Err(e) => {
*rep.reject_reasons.entry(format!("provider: {e}")).or_default() += 1;
continue;
}
};
let payload = turn
.tool_uses
.first()
.map(|(_, _, v)| v.clone())
.or_else(|| crate::vocabulary::extract_json(&turn.text));
let Some(v) = payload else {
*rep.reject_reasons.entry("no structured output".into()).or_default() += 1;
continue;
};
let (kept, rejects) = examples_from_proposal(spec, &v, case);
*rep.per_case.entry(format!("{case:?}").to_lowercase()).or_default() += kept.len();
rep.kept += kept.len();
rep.rejected += rejects.len();
for r in rejects {
let key = match r {
Reject::SurfaceNotFound(_) => "surface_not_found",
Reject::UndeclaredFacet(_) => "undeclared_facet",
Reject::UndeclaredRelation(_) => "undeclared_relation",
Reject::TypeMismatch { .. } => "type_mismatch",
Reject::OverlappingSpans => "overlapping_spans",
Reject::BadRelationIndex => "bad_relation_index",
Reject::NoSpans => "no_spans",
};
*rep.reject_reasons.entry(key.into()).or_default() += 1;
}
out.extend(kept);
}
}
(out, rep)
}
/// Serialise a dataset as JSONL (one example per line).
pub fn to_jsonl(examples: &[TaggerExample]) -> String {
examples.iter().filter_map(|e| serde_json::to_string(e).ok()).map(|l| l + "\n").collect()
}
// ── corpus-grounded generation (the `tune_ontology.py` correction) ──────────────────────────────
//
// Inventing sentences teaches the heads a vocabulary the corpus does not have — the reference names this
// exact failure ("the heads never learned the corpus's real vocabulary"). Grounded generation instead hands
// the model REAL passages and asks it to label only what is present. `validate` then enforces that
// mechanically: a surface that is not a substring of the passage is rejected, so "do not invent" is a
// checked property rather than an instruction we hope was followed.
/// Prose passages worth labelling, pulled from sampled documents. Field-label lines ("**Id:** 1a2b") carry
/// no relational language, so only sentence-like segments are kept.
pub fn passages_from_docs(docs: &[String], min_words: usize, max_chars: usize) -> Vec<String> {
let mut out: Vec<String> = Vec::new();
for doc in docs {
for raw in doc.split(['\n', '\r']) {
let line = raw.trim().trim_start_matches(['-', '*', '#', ' ']).trim();
// Strip a leading field label, keeping the prose after it. The bullet/asterisk trim above has
// already eaten the opening `**`, so what remains looks like `Abstract:** rest` or `Abstract: rest`;
// a short prefix ending in a colon is a label, not prose.
let line = match line.find(':') {
Some(i) if i <= 40 => line[i + 1..].trim_start_matches(['*', ' ']).trim(),
_ => line,
};
let words = line.split_whitespace().count();
if words < min_words || line.len() < 40 {
continue;
}
// must read like prose: contains a verb-ish lowercase word and isn't mostly identifiers
let alpha = line.chars().filter(|c| c.is_alphabetic()).count();
if alpha * 2 < line.len() {
continue;
}
out.push(line.chars().take(max_chars).collect());
}
}
out
}
pub const GROUNDED_SYSTEM: &str = "You label real corpus passages for a span-tagging and relation-extraction model. For each passage, extract ONLY terms that ACTUALLY APPEAR in that passage — copy each 'surface' VERBATIM as an exact substring. Do NOT invent entities, and do not paraphrase. Assign each span one of the given facets. Add relations only where the passage genuinely asserts one, referencing spans by their exact surfaces and respecting the declared head/tail facet types. If a passage contains nothing relevant, return no spans for it. Mark negated=true or hedged=true on spans whose assertion the passage negates or hedges, and label the negation/hedge cue itself as a span with facet \"state\".";
/// Schema for labelling a batch of real passages (indexed, so replies map back to their source text).
pub fn grounded_schema() -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": { "passages": { "type": "array", "items": { "type": "object", "properties": {
"index": {"type": "integer"},
"spans": {"type": "array", "items": {"type": "object", "properties": {
"surface": {"type": "string"}, "facet": {"type": "string"},
"negated": {"type": "boolean"}, "hedged": {"type": "boolean"}
}, "required": ["surface", "facet"]}},
"relations": {"type": "array", "items": {"type": "object", "properties": {
"head_surface": {"type": "string"}, "tail_surface": {"type": "string"}, "name": {"type": "string"}
}, "required": ["head_surface", "tail_surface", "name"]}}
}, "required": ["index", "spans"] } } },
"required": ["passages"]
})
}
/// Descriptions for the corpus-independent dimensions, included in every grounded request. Without these
/// in the menu the model never labels them, leaving Head A's `GEO`/`TIME`/`QTY`/`STATE` classes with no
/// training signal at all — present in the label space but unable to fire.
pub const STRUCTURAL_MENU: &[(&str, &str)] = &[
("geo", "a place: country, region, city, or named locale"),
("time", "a time expression: year, quarter, month, date, or named period"),
("qty", "a quantity with its unit: durations, counts with units, sizes, temperatures, currency amounts"),
("state", "the negation or hedging cue itself (\"does not\", \"may\", \"reportedly\")"),
];
/// Prompt for one batch of real passages.
pub fn grounded_prompt(spec: &VocabularySpace, passages: &[String]) -> String {
let mut facets: Vec<String> = spec.taggable_facets().iter().map(|f| format!("{} — {}", f.name, f.description)).collect();
facets.extend(STRUCTURAL_MENU.iter().map(|(n, d)| format!("{n} — {d}")));
let rels: Vec<String> = spec.relation_facets.iter().map(|r| format!("{} ({} → {})", r.name, r.head, r.tail)).collect();
let body: Vec<String> = passages.iter().enumerate().map(|(i, p)| format!("[{i}] {p}")).collect();
format!(
"Entity facets:\n{}\n\nRelations:\n{}\n\nPassages:\n{}",
facets.join("\n"),
rels.join("\n"),
body.join("\n\n")
)
}
/// Turn a grounded reply into examples, aligning surfaces against the ORIGINAL passage text.
pub fn examples_from_grounded(
spec: &VocabularySpace,
passages: &[String],
v: &serde_json::Value,
) -> (Vec<TaggerExample>, Vec<Reject>) {
let mut kept = Vec::new();
let mut rejects = Vec::new();
let Some(items) = v.get("passages").and_then(|x| x.as_array()) else { return (kept, rejects) };
for it in items {
let Some(idx) = it.get("index").and_then(|x| x.as_u64()).map(|n| n as usize) else { continue };
let Some(text) = passages.get(idx) else { continue };
// reuse the invented-path builder by re-shaping this entry into its schema
let one = serde_json::json!({ "examples": [{
"text": text,
"spans": it.get("spans").cloned().unwrap_or(serde_json::json!([])),
"relations": it.get("relations").cloned().unwrap_or(serde_json::json!([])),
}]});
let (k, r) = examples_from_proposal(spec, &one, Case::Normal);
kept.extend(k);
rejects.extend(r);
}
(kept, rejects)
}
/// Mine a gazetteer from labelled data: every span is already a `(surface, facet)` pair verified to occur
/// in real text, so the high-resolution whole-entity vocabulary falls out of step 1 with no extra model
/// calls. Multi-word surfaces only — single tokens are what the tagger and SPLADE tiers already cover, and
/// the gazetteer exists precisely to keep multi-word entities from shattering.
pub fn mine_gazetteer(
spec: &VocabularySpace,
examples: &[TaggerExample],
min_count: usize,
) -> Vec<crate::vocabulary::GazEntry> {
use std::collections::BTreeMap;
// (normalised surface, facet) → (count, best original casing)
let mut seen: BTreeMap<(String, String), (usize, String)> = BTreeMap::new();
for ex in examples {
for sp in &ex.spans {
if sp.facet == "state" || STRUCTURAL_KINDS.contains(&sp.facet.to_uppercase().as_str()) {
continue; // loci/quantities are normalised deterministically, not gazetteered
}
let surface = sp.surface.trim();
if surface.split_whitespace().count() < 2 || surface.len() < 4 {
continue;
}
// A determiner-led phrase ("the feature", "its workspace") names no specific entity — the
// gazetteer is for high-resolution whole entities, so generic references are noise in it.
const LEADING_GENERIC: &[&str] = &["the", "a", "an", "this", "that", "these", "those", "its", "their", "our", "your", "his", "her", "such", "any", "each"];
let first = surface.split_whitespace().next().unwrap_or("").to_lowercase();
if LEADING_GENERIC.contains(&first.as_str()) {
continue;
}
let key = (surface.to_lowercase(), sp.facet.clone());
let e = seen.entry(key).or_insert((0, surface.to_string()));
e.0 += 1;
}
}
seen.into_iter()
.filter(|(_, (n, _))| *n >= min_count)
.filter_map(|((_, facet), (_, surface))| {
// facet-qualified, hierarchical token so wildcards reach it
spec.has_entity_facet(&facet).then(|| crate::vocabulary::GazEntry {
token: spec.entity_uri(&facet, &surface),
surface,
})
})
.collect()
}
/// Label real corpus passages in batches. Empty results are normal — most passages in a structured corpus
/// carry no relational language — so the report tracks how many passages yielded anything.
#[cfg(feature = "agent")]
pub async fn generate_grounded(
provider: &dyn crate::agent::provider::LlmProvider,
spec: &VocabularySpace,
passages: &[String],
batch: usize,
) -> (Vec<TaggerExample>, GenReport) {
use crate::agent::types::{Msg, ToolSpec};
let tools = vec![ToolSpec {
name: "emit_labels".into(),
description: "Emit span/relation labels for each passage.".into(),
schema: grounded_schema(),
}];
let mut out: Vec<TaggerExample> = Vec::new();
let mut rep = GenReport::default();
for chunk in passages.chunks(batch.max(1)) {
let prompt = grounded_prompt(spec, chunk);
let turn = match provider.chat(GROUNDED_SYSTEM, &[Msg::user_text(prompt)], &tools).await {
Ok(t) => t,
Err(e) => {
*rep.reject_reasons.entry(format!("provider: {e}")).or_default() += 1;
continue;
}
};
let payload = turn
.tool_uses
.first()
.map(|(_, _, v)| v.clone())
.or_else(|| crate::vocabulary::extract_json(&turn.text));
let Some(v) = payload else {
*rep.reject_reasons.entry("no structured output".into()).or_default() += 1;
continue;
};
let (kept, rejects) = examples_from_grounded(spec, chunk, &v);
*rep.per_case.entry("grounded".into()).or_default() += kept.len();
rep.kept += kept.len();
rep.rejected += rejects.len();
for r in rejects {
let key = match r {
Reject::SurfaceNotFound(_) => "surface_not_found",
Reject::UndeclaredFacet(_) => "undeclared_facet",
Reject::UndeclaredRelation(_) => "undeclared_relation",
Reject::TypeMismatch { .. } => "type_mismatch",
Reject::OverlappingSpans => "overlapping_spans",
Reject::BadRelationIndex => "bad_relation_index",
Reject::NoSpans => "no_spans",
};
*rep.reject_reasons.entry(key.into()).or_default() += 1;
}
out.extend(kept);
}
(out, rep)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::vocabulary::{EntityFacet, RelationFacet};
fn spec() -> VocabularySpace {
VocabularySpace {
version: 1,
corpus: "defence".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 head_spaces_derive_from_spec() {
let s = spec();
let a = head_a_labels(&s);
assert_eq!(a[0], "O");
assert!(a.contains(&"B-ORG".to_string()) && a.contains(&"I-SYSTEM".to_string()));
// structural kinds always present
assert!(a.contains(&"B-QTY".to_string()) && a.contains(&"B-TIME".to_string()));
assert_eq!(a.len(), 1 + 2 * (2 + STRUCTURAL_KINDS.len()));
assert_eq!(head_c_labels(&s), vec!["none", "develops"]);
// the type mask makes an invalid pairing unrepresentable
let m = pair_mask(&s);
assert!(m.contains(&("org".into(), "system".into(), "develops".into())));
assert!(!m.contains(&("system".into(), "org".into(), "develops".into())));
}
#[test]
fn alignment_finds_offsets_and_avoids_overlap() {
let t = "Boeing develops the MQ-28, and Boeing also funds it.";
let a = align(t, "Boeing", &[]).unwrap();
assert_eq!(&t[a.0..a.1], "Boeing");
// second mention when the first is taken
let b = align(t, "Boeing", &[a]).unwrap();
assert!(b.0 > a.0);
assert_eq!(&t[b.0..b.1], "Boeing");
// case-insensitive, offsets into original
let c = align(t, "mq-28", &[]).unwrap();
assert_eq!(&t[c.0..c.1], "MQ-28");
assert!(align(t, "Airbus", &[]).is_none());
}
#[test]
fn align_offsets_survive_non_ascii() {
// an em dash before the target: offsets from a lowercased copy drift here
let t = "Amazon Connect — the customer's staff use it.";
let (s0, e0) = align(t, "customer", &[]).unwrap();
assert_eq!(&t[s0..e0], "customer", "offsets must index the ORIGINAL text");
let (s1, e1) = align(t, "Amazon Connect", &[]).unwrap();
assert_eq!(&t[s1..e1], "Amazon Connect");
// curly apostrophe + case differences
let t2 = "The CUSTOMER’S centre — a medical centre — closed.";
let (s2, e2) = align(t2, "medical centre", &[]).unwrap();
assert_eq!(&t2[s2..e2], "medical centre");
}
#[test]
fn align_respects_word_boundaries() {
// regression: "it" must NOT match inside "submitted" (this silently mislabelled real generated data)
let t = "Boeing submitted documentation, and they are developing it for airframes.";
let (s, e) = align(t, "it", &[]).unwrap();
assert_eq!(&t[s..e], "it");
assert!(s > 50, "must find the standalone pronoun, not the one inside 'submitted' (got offset {s})");
// "the aircraft" inside a longer phrase is fine; punctuated surfaces still align
let t2 = "They fly the MQ-28, a loyal wingman.";
assert_eq!(align(t2, "MQ-28", &[]).map(|(a, b)| &t2[a..b]), Some("MQ-28"));
// a surface that only occurs as a sub-word is correctly rejected
assert!(align("Documentation submitted.", "it", &[]).is_none());
}
#[test]
fn validation_rejects_bad_examples() {
let s = spec();
let mut 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,
};
assert!(validate(&s, &ex).is_ok());
// reversed arguments violate the declared head/tail types
ex.relations = vec![RelationLabel { head: 1, tail: 0, name: "develops".into() }];
assert!(matches!(validate(&s, &ex), Err(Reject::TypeMismatch { .. })));
// undeclared facet
ex.relations.clear();
ex.spans[0].facet = "gene".into();
assert_eq!(validate(&s, &ex), Err(Reject::UndeclaredFacet("gene".into())));
// offsets that don't match the surface
ex.spans[0].facet = "org".into();
ex.spans[0].end = 5;
assert!(matches!(validate(&s, &ex), Err(Reject::SurfaceNotFound(_))));
}
#[test]
fn proposal_to_examples_aligns_and_filters() {
let s = spec();
let v = serde_json::json!({"examples": [
// good
{"text": "Boeing develops the MQ-28 Ghost Bat.",
"spans": [{"surface":"Boeing","facet":"org"},{"surface":"MQ-28 Ghost Bat","facet":"system"}],
"relations": [{"head_surface":"Boeing","tail_surface":"MQ-28 Ghost Bat","name":"develops"}]},
// surface not in text → rejected
{"text": "Airbus builds jets.", "spans": [{"surface":"Boeing","facet":"org"}]},
// type-invalid relation → rejected
{"text": "The MQ-28 develops Boeing.",
"spans": [{"surface":"MQ-28","facet":"system"},{"surface":"Boeing","facet":"org"}],
"relations": [{"head_surface":"MQ-28","tail_surface":"Boeing","name":"develops"}]}
]});
let (kept, rejects) = examples_from_proposal(&s, &v, Case::Normal);
assert_eq!(kept.len(), 1);
assert_eq!(kept[0].spans.len(), 2);
assert_eq!(kept[0].relations[0].name, "develops");
assert_eq!(rejects.len(), 2);
assert!(rejects.iter().any(|r| matches!(r, Reject::SurfaceNotFound(_))));
assert!(rejects.iter().any(|r| matches!(r, Reject::TypeMismatch { .. })));
}
#[test]
fn bio_and_epistemic_projection() {
let s = spec();
let ex = TaggerExample {
text: "Boeing develops MQ-28".into(),
spans: vec![
LabeledSpan { start: 0, end: 6, facet: "org".into(), surface: "Boeing".into(), negated: true, hedged: false },
LabeledSpan { start: 16, end: 21, facet: "system".into(), surface: "MQ-28".into(), negated: false, hedged: false },
],
relations: vec![],
case: Case::Negated,
};
// token offsets: [CLS] Boeing develops MQ - 28 [SEP]
let offsets = [(0, 0), (0, 6), (7, 15), (16, 18), (18, 19), (19, 21), (0, 0)];
let bio = to_bio(&s, &ex.spans, &offsets);
let labels = head_a_labels(&s);
assert_eq!(bio[0], -100); // special
assert_eq!(labels[bio[1] as usize], "B-ORG");
assert_eq!(bio[2], 0); // "develops" is O (the REL span isn't labelled in this example)
assert_eq!(labels[bio[3] as usize], "B-SYSTEM"); // first token of the span
assert_eq!(labels[bio[4] as usize], "I-SYSTEM"); // continuation
let ep = to_epistemic(&ex.spans, &offsets);
assert_eq!(ep[1], 3); // negated
assert_eq!(ep[3], 0); // asserted
}
#[test]
fn passage_extraction_keeps_prose_and_drops_field_labels() {
let doc = "# Source Record\n- **Id:** CN4082207-f074\n- **Abstract:** The proposed effort develops a compact lithium-ion battery module for unmanned maritime platforms.\n- **Status:** active\n";
let p = passages_from_docs(&[doc.to_string()], 8, 400);
assert_eq!(p.len(), 1, "only the prose abstract qualifies: {p:?}");
assert!(p[0].starts_with("The proposed effort develops"), "label stripped, prose kept: {:?}", p[0]);
// identifier-heavy and short lines are excluded
assert!(!p.iter().any(|x| x.contains("CN4082207")));
assert!(!p.iter().any(|x| x.contains("active")));
}
#[test]
fn grounded_replies_align_against_the_real_passage() {
let s = spec();
let passages = vec!["Boeing develops the MQ-28 Ghost Bat for the Royal Australian Air Force.".to_string()];
let v = serde_json::json!({"passages": [{
"index": 0,
"spans": [{"surface":"Boeing","facet":"org"},{"surface":"MQ-28 Ghost Bat","facet":"system"}],
"relations": [{"head_surface":"Boeing","tail_surface":"MQ-28 Ghost Bat","name":"develops"}]
}]});
let (kept, rejects) = examples_from_grounded(&s, &passages, &v);
assert_eq!(kept.len(), 1);
assert!(rejects.is_empty());
assert_eq!(&kept[0].text[kept[0].spans[0].start..kept[0].spans[0].end], "Boeing");
// an INVENTED surface is mechanically rejected — "do not invent" is enforced, not trusted
let bad = serde_json::json!({"passages": [{
"index": 0, "spans": [{"surface":"Lockheed Martin","facet":"org"}]
}]});
let (k2, r2) = examples_from_grounded(&s, &passages, &bad);
assert!(k2.is_empty());
assert!(matches!(r2.first(), Some(Reject::SurfaceNotFound(_))));
}
#[test]
fn overloaded_relation_names_accept_every_declared_signature() {
// `has_type` legitimately applies to two different head facets
let mut s = spec();
s.entity_facets.push(EntityFacet { name: "move".into(), parent: None, description: "moves".into(), examples: vec![], structural: false });
s.entity_facets.push(EntityFacet { name: "kind".into(), parent: None, description: "types".into(), examples: vec![], structural: false });
s.relation_facets.push(RelationFacet { name: "has_type".into(), head: "system".into(), tail: "kind".into() });
s.relation_facets.push(RelationFacet { name: "has_type".into(), head: "move".into(), tail: "kind".into() });
let mk = |hf: &str, tf: &str| TaggerExample {
text: "Iron Head is a steel move used by Metagross.".into(),
spans: vec![
LabeledSpan { start: 0, end: 9, facet: hf.into(), surface: "Iron Head".into(), negated: false, hedged: false },
LabeledSpan { start: 15, end: 20, facet: tf.into(), surface: "steel".into(), negated: false, hedged: false },
],
relations: vec![RelationLabel { head: 0, tail: 1, name: "has_type".into() }],
case: Case::Normal,
};
// both declared signatures must validate
assert!(validate(&s, &mk("move", "kind")).is_ok(), "move → kind is declared");
assert!(validate(&s, &mk("system", "kind")).is_ok(), "system → kind is also declared");
// an undeclared pairing still fails
assert!(matches!(validate(&s, &mk("kind", "move")), Err(Reject::TypeMismatch { .. })));
}
#[test]
fn gazetteer_mining_keeps_multiword_entities_only() {
let s = spec();
let mk = |text: &str, spans: Vec<(usize, usize, &str)>| TaggerExample {
text: text.into(),
spans: spans
.into_iter()
.map(|(a, b, f)| LabeledSpan { start: a, end: b, facet: f.into(), surface: text[a..b].into(), negated: false, hedged: false })
.collect(),
relations: vec![],
case: Case::Normal,
};
let t = "Amazon Connect Contact Lens helps Boeing in Sydney.";
let examples = vec![
mk(t, vec![(0, 27, "org"), (34, 40, "org"), (44, 50, "geo")]),
mk(t, vec![(0, 27, "org")]), // seen twice → clears min_count
];
let g = mine_gazetteer(&s, &examples, 2);
let surfaces: Vec<&str> = g.iter().map(|e| e.surface.as_str()).collect();
assert_eq!(surfaces, vec!["Amazon Connect Contact Lens"], "multi-word, repeated, entity-facet only");
assert_eq!(g[0].token, "org/amazon-connect-contact-lens", "facet-qualified for wildcards");
// single-word entities and loci are excluded
assert!(!surfaces.contains(&"Boeing"));
assert!(!surfaces.iter().any(|x| *x == "Sydney"));
// determiner-led generic references are not entities
let generic = vec![mk("the feature helps the feature", vec![(0, 11, "org")]), mk("the feature helps the feature", vec![(0, 11, "org")])];
assert!(mine_gazetteer(&s, &generic, 2).is_empty(), "determiner-led phrases must be filtered");
// below min_count → nothing
assert!(mine_gazetteer(&s, &examples[1..], 2).is_empty());
}
#[test]
fn structural_dimensions_appear_in_the_grounded_menu() {
let s = spec();
let p = grounded_prompt(&s, &["Boeing shipped 12 units in Q3 2026 to Sydney.".to_string()]);
for dim in ["geo —", "time —", "qty —", "state —"] {
assert!(p.contains(dim), "grounded menu must offer {dim}: {p}");
}
assert!(p.contains("org —"), "spec facets still offered");
}
#[test]
fn plan_guarantees_hard_case_coverage() {
let p = GenPlan::default();
assert_eq!(p.total_per_relation(), 12);
let cases = p.cases();
for hard in [Case::Coref, Case::Hedged, Case::Negated, Case::Adversarial] {
assert!(cases.iter().any(|(c, n)| *c == hard && *n > 0), "{hard:?} must be planned");
}
}
}