uni_xervo/traits/nlp.rs
1//! Structured natural-language analysis types and traits.
2//!
3//! The trait is designed around multi-head cascade models (e.g.
4//! `dragonscale-ai/kniv-deberta-nlp-base-en-xsmall`) that produce
5//! token-level POS / NER / DEP tags, sentence segmentation, optional SRL
6//! frames, and optional dialog-act classification in a single forward pass.
7
8use crate::error::Result;
9use crate::traits::ModelInfo;
10use async_trait::async_trait;
11use bitflags::bitflags;
12
13bitflags! {
14 /// Selects which NLP heads a caller wants populated.
15 ///
16 /// Implementations that support a head only do its post-processing
17 /// (label decoding, dependency tree reconstruction, SRL frame
18 /// assembly) when its flag is set — this saves measurable wall time
19 /// on hot ingest paths where, for example, only NER is needed.
20 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
21 pub struct NlpTasks: u32 {
22 /// Part-of-speech tagging.
23 const POS = 1 << 0;
24 /// Named entity recognition.
25 const NER = 1 << 1;
26 /// Dependency parsing.
27 const DEP = 1 << 2;
28 /// Semantic role labeling.
29 const SRL = 1 << 3;
30 /// Sentence-level classification (e.g. dialog acts).
31 const CLS = 1 << 4;
32 /// Convenience: all currently-defined heads.
33 const ALL = Self::POS.bits() | Self::NER.bits() | Self::DEP.bits()
34 | Self::SRL.bits() | Self::CLS.bits();
35 }
36}
37
38/// One NLP analysis request: a text plus the heads to populate.
39#[derive(Debug, Clone)]
40pub struct NlpRequest<'a> {
41 /// Input text. Borrowed so batch callers don't have to clone.
42 pub text: &'a str,
43 /// Which heads to populate. Unset flags result in empty fields on the
44 /// corresponding [`NlpResult`] entries.
45 pub tasks: NlpTasks,
46}
47
48/// Structured output of an NLP analysis call.
49///
50/// Field population depends on the [`NlpTasks`] flags in the request and on
51/// the model's [`NlpModel::supported_tasks`]. Unrequested or unsupported
52/// fields are empty (`Vec::new()`).
53///
54/// This struct is `#[non_exhaustive]`: construct it via [`NlpResult::default`]
55/// and assign fields, so future heads can be added without breaking callers.
56#[derive(Debug, Clone, Default)]
57#[non_exhaustive]
58pub struct NlpResult {
59 /// Token sequence with optional POS / NER / DEP annotations.
60 pub tokens: Vec<NlpToken>,
61 /// Sentence boundaries (always populated when any head is requested).
62 pub sentences: Vec<NlpSentence>,
63 /// SRL frames. Empty unless [`NlpTasks::SRL`] was requested and supported.
64 pub frames: Vec<SrlFrame>,
65 /// Sentence-level classifications. Empty unless [`NlpTasks::CLS`] was
66 /// requested and supported.
67 pub speech_acts: Vec<SpeechAct>,
68 /// Merged named-entity spans, a BIO-collapsed view of the per-token
69 /// [`NlpToken::ner`] tags. Empty unless [`NlpTasks::NER`] was requested
70 /// and supported. See [`NerEntity`].
71 pub entities: Vec<NerEntity>,
72}
73
74/// One token in an [`NlpResult`].
75///
76/// Tokens are subword units; several may belong to one word (see
77/// [`word_index`](NlpToken::word_index)). Index-based references elsewhere
78/// (e.g. [`DepLink::head`], [`SrlRole::span`]) are always *token* indices into
79/// [`NlpResult::tokens`], not word indices.
80#[derive(Debug, Clone, Default)]
81#[non_exhaustive]
82pub struct NlpToken {
83 /// Surface form of the token.
84 pub text: String,
85 /// UTF-8 byte offset of the token start in the original input text.
86 pub start: usize,
87 /// UTF-8 byte offset of the token end (exclusive).
88 pub end: usize,
89 /// POS tag (typically Universal Dependencies tagset). `None` if POS not requested.
90 pub pos: Option<String>,
91 /// Named-entity type (e.g. `"PERSON"`, `"ORG"`). `None` outside any entity span.
92 pub ner: Option<String>,
93 /// Dependency-parse head and relation label. `None` if DEP not requested.
94 pub dep: Option<DepLink>,
95 /// Index of the word this subword token belongs to.
96 ///
97 /// Dense and monotonically non-decreasing over [`NlpResult::tokens`]:
98 /// consecutive tokens that form one word share a `word_index`, and each new
99 /// word increments it by one starting at `0`. Lets consumers regroup
100 /// subwords into words without parsing tokenizer metaspace markers.
101 pub word_index: usize,
102}
103
104/// One sentence in an [`NlpResult`].
105#[derive(Debug, Clone, Default)]
106#[non_exhaustive]
107pub struct NlpSentence {
108 /// Inclusive `[first, last]` token indices into [`NlpResult::tokens`].
109 pub token_range: (usize, usize),
110 /// UTF-8 byte offset of the sentence start in the original text.
111 pub start: usize,
112 /// UTF-8 byte offset of the sentence end (exclusive).
113 pub end: usize,
114}
115
116/// A dependency-parse arc attached to an [`NlpToken`].
117#[derive(Debug, Clone, Default, PartialEq, Eq)]
118#[non_exhaustive]
119pub struct DepLink {
120 /// Syntactic head as a 0-based index into [`NlpResult::tokens`].
121 ///
122 /// `Some(i)` points at the head token; `None` means this token attaches to
123 /// the sentence root (it has no head). The index is global across the whole
124 /// [`NlpResult`], not chunk- or sentence-local.
125 pub head: Option<usize>,
126 /// Relation label (e.g. `"nsubj"`, `"obj"`, `"amod"`).
127 pub relation: String,
128}
129
130/// One semantic-role-labeling frame: a predicate and its argument spans.
131#[derive(Debug, Clone, Default)]
132#[non_exhaustive]
133pub struct SrlFrame {
134 /// 0-based index into [`NlpResult::tokens`] of the predicate token.
135 ///
136 /// The predicate token is never itself included in [`roles`](SrlFrame::roles).
137 pub predicate_token: usize,
138 /// Predicate sense identifier when known (e.g. `"buy.01"`).
139 pub predicate_sense: Option<String>,
140 /// Argument spans, excluding the predicate token.
141 pub roles: Vec<SrlRole>,
142}
143
144/// One argument span of an [`SrlFrame`].
145#[derive(Debug, Clone, Default, PartialEq, Eq)]
146#[non_exhaustive]
147pub struct SrlRole {
148 /// Inclusive `[first, last]` 0-based token indices into [`NlpResult::tokens`].
149 ///
150 /// Both endpoints are inclusive; a single-token argument has `first == last`.
151 pub span: (usize, usize),
152 /// Role label (e.g. `"ARG0"`, `"ARG1"`, `"ARGM-TMP"`).
153 pub label: String,
154}
155
156/// One sentence-level classification result.
157#[derive(Debug, Clone, Default)]
158#[non_exhaustive]
159pub struct SpeechAct {
160 /// Index into [`NlpResult::sentences`].
161 pub sentence_index: usize,
162 /// Class label (e.g. `"STATEMENT"`, `"QUESTION"`, `"GREETING"`).
163 ///
164 /// Equals the label of the argmax of [`scores`](SpeechAct::scores).
165 pub label: String,
166 /// Provider-reported confidence in `[0.0, 1.0]`.
167 ///
168 /// Equals `scores[argmax]` — the probability of the winning [`label`](SpeechAct::label).
169 pub confidence: f32,
170 /// Full per-class softmax over the model's dialog-act vocabulary.
171 ///
172 /// Parallel to [`NlpLabelMaps::cls`]: `scores[i]` is the probability of the
173 /// class named `cls[i]`, and the values sum to ~1.0. Empty when the model
174 /// does not expose a distribution. Lets consumers do thresholded
175 /// multi-label gating instead of relying on the top-1 [`label`](SpeechAct::label).
176 pub scores: Vec<f32>,
177}
178
179/// One merged named-entity span, a BIO-collapsed view of per-token NER tags.
180///
181/// Produced by collapsing the per-token [`NlpToken::ner`] BIO tags: a `B-` (or
182/// orphan `I-`) tag opens a span, matching `I-` tags extend it, and an `O` tag
183/// or a label change closes it. Spans never include the `"O"` outside-tag.
184#[derive(Debug, Clone, Default, PartialEq, Eq)]
185#[non_exhaustive]
186pub struct NerEntity {
187 /// Surface text of the entity, spanning the original input bytes.
188 pub text: String,
189 /// Entity type without the BIO prefix (e.g. `"PERSON"`, `"GPE"`).
190 pub label: String,
191 /// Inclusive `[first, last]` 0-based token indices into [`NlpResult::tokens`].
192 pub token_span: (usize, usize),
193 /// Half-open `[start, end)` UTF-8 byte offsets into the original input text.
194 pub char_span: (usize, usize),
195}
196
197/// The label vocabularies a model decodes against, one list per head.
198///
199/// Each vector is indexed by the model's internal class id, so a decoded class
200/// id `k` for a head maps to `labels.<head>[k]`. In particular
201/// [`cls`](NlpLabelMaps::cls) is parallel to [`SpeechAct::scores`]. Exposing
202/// these lets consumers map labels back to ids without embedding a parallel
203/// copy of the model's `label_maps.json`.
204#[derive(Debug, Clone, Default)]
205#[non_exhaustive]
206pub struct NlpLabelMaps {
207 /// Part-of-speech tag vocabulary.
208 pub pos: Vec<String>,
209 /// Named-entity BIO tag vocabulary.
210 pub ner: Vec<String>,
211 /// Dependency relation vocabulary.
212 pub deprel: Vec<String>,
213 /// Semantic-role BIO tag vocabulary.
214 pub srl: Vec<String>,
215 /// Dialog-act / sentence-classification vocabulary, parallel to
216 /// [`SpeechAct::scores`].
217 pub cls: Vec<String>,
218}
219
220/// A multi-head structured-NLP model.
221///
222/// # Canonical default model
223///
224/// The reference implementation is
225/// [`dragonscale-ai/kniv-deberta-nlp-base-en-xsmall`](https://huggingface.co/dragonscale-ai/kniv-deberta-nlp-base-en-xsmall):
226/// a DeBERTa-v3-xsmall multi-head cascade producing POS / NER / DEP / SRL / CLS
227/// in a single forward pass. It will ship as the canonical `nlp/default`
228/// catalog alias once the `local/onnx` provider gains an `NlpModel`
229/// implementation in a follow-up release.
230///
231/// Tagsets used by the reference model:
232/// - POS: Universal Dependencies English EWT (17 classes)
233/// - NER: OntoNotes (18 entity types)
234/// - DEP: Universal Dependencies English EWT
235/// - SRL: PropBank (42 classes)
236/// - CLS: 8 dialog acts (internal dataset)
237///
238/// Language: English only. Maximum input: 128 tokens — provider-side
239/// chunking is the implementation's responsibility; the trait contract
240/// keeps token offsets stable in whole-document UTF-8 byte coordinates
241/// regardless of internal chunking.
242#[async_trait]
243pub trait NlpModel: ModelInfo {
244 /// Analyze a batch of texts, returning one [`NlpResult`] per request.
245 ///
246 /// # Errors
247 /// Returns an error if a requested task is not in
248 /// [`supported_tasks`](NlpModel::supported_tasks), or if the provider
249 /// fails internally.
250 async fn analyze(&self, requests: Vec<NlpRequest<'_>>) -> Result<Vec<NlpResult>>;
251
252 /// Heads this model can populate. Callers may request a subset via
253 /// [`NlpRequest::tasks`].
254 fn supported_tasks(&self) -> NlpTasks;
255
256 /// The label vocabularies this model decodes against, if it exposes them.
257 ///
258 /// Returns `Some` for models with a fixed, introspectable label set (e.g.
259 /// the ONNX cascade), letting callers map labels to class ids and read
260 /// [`SpeechAct::scores`] against [`NlpLabelMaps::cls`]. The default is
261 /// `None` for models that do not expose vocabularies (e.g. remote or mock).
262 fn label_maps(&self) -> Option<&NlpLabelMaps> {
263 None
264 }
265
266 /// Optional warmup hook. The default is a no-op.
267 ///
268 /// # Errors
269 /// Returns an error if the underlying model fails to initialize.
270 async fn warmup(&self) -> Result<()> {
271 Ok(())
272 }
273}
274
275#[cfg(test)]
276mod tests {
277 use super::*;
278 use std::sync::Arc;
279
280 #[test]
281 fn nlp_model_is_dyn_safe() {
282 fn _accept(_: Arc<dyn NlpModel>) {}
283 }
284
285 #[test]
286 fn nlp_tasks_all_includes_every_head() {
287 let all = NlpTasks::ALL;
288 for flag in [
289 NlpTasks::POS,
290 NlpTasks::NER,
291 NlpTasks::DEP,
292 NlpTasks::SRL,
293 NlpTasks::CLS,
294 ] {
295 assert!(all.contains(flag), "{flag:?} missing from NlpTasks::ALL");
296 }
297 }
298}