sie_sdk/types/
encode.rs

1//! Results of `/v1/encode`, `/v1/score` and `/v1/extract`.
2
3// Wire-mirror types: field names are the API contract itself, and the ones whose
4// meaning is not obvious carry their own doc comment.
5#![allow(missing_docs)]
6
7use half::f16;
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10
11use super::{RequestMetadata, TimingInfo};
12
13/// A sparse embedding: term ids paired with their weights.
14#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
15pub struct SparseVector {
16    pub indices: Vec<u32>,
17    pub values: Vec<f32>,
18}
19
20impl SparseVector {
21    /// The `{term_id: weight}` form most vector stores expect.
22    pub fn to_map(&self) -> std::collections::HashMap<u32, f32> {
23        self.indices
24            .iter()
25            .copied()
26            .zip(self.values.iter().copied())
27            .collect()
28    }
29
30    /// How many terms the vector carries.
31    pub fn len(&self) -> usize {
32        self.indices.len().min(self.values.len())
33    }
34
35    /// Whether the vector carries no terms.
36    pub fn is_empty(&self) -> bool {
37        self.len() == 0
38    }
39}
40
41/// A token-level embedding matrix, in the precision the model produced.
42///
43/// Late-interaction models emit `f16` by default; the SDK keeps that precision rather than
44/// silently widening, because [`crate::scoring::maxsim`] has to control exactly where the
45/// widening happens.
46#[derive(Debug, Clone, PartialEq)]
47pub enum Multivector {
48    F16(Vec<Vec<f16>>),
49    F32(Vec<Vec<f32>>),
50}
51
52impl Multivector {
53    /// Number of token vectors.
54    pub fn len(&self) -> usize {
55        match self {
56            Self::F16(rows) => rows.len(),
57            Self::F32(rows) => rows.len(),
58        }
59    }
60
61    /// Whether there are no token vectors.
62    pub fn is_empty(&self) -> bool {
63        self.len() == 0
64    }
65
66    /// Width of each token vector.
67    pub fn dims(&self) -> usize {
68        match self {
69            Self::F16(rows) => rows.first().map_or(0, Vec::len),
70            Self::F32(rows) => rows.first().map_or(0, Vec::len),
71        }
72    }
73
74    /// Widen to `f32` rows, copying when the wire precision was `f16`.
75    pub fn to_f32(&self) -> Vec<Vec<f32>> {
76        match self {
77            Self::F16(rows) => rows
78                .iter()
79                .map(|row| row.iter().map(|value| value.to_f32()).collect())
80                .collect(),
81            Self::F32(rows) => rows.clone(),
82        }
83    }
84}
85
86/// One encoded item.
87#[derive(Debug, Clone, Default, PartialEq)]
88pub struct EncodeResult {
89    pub model: Option<String>,
90    /// Echoed from the request item, when it carried an id.
91    pub id: Option<String>,
92    pub dense: Option<Vec<f32>>,
93    pub sparse: Option<SparseVector>,
94    pub multivector: Option<Multivector>,
95    pub timing: Option<TimingInfo>,
96    pub request: Option<RequestMetadata>,
97}
98
99impl EncodeResult {
100    /// The dense embedding, or an error naming what is missing.
101    ///
102    /// Use the `dense` field directly when its absence is expected.
103    pub fn require_dense(&self) -> crate::error::Result<&[f32]> {
104        self.dense
105            .as_deref()
106            .ok_or_else(|| crate::error::Error::decode("encode result has no dense embedding"))
107    }
108
109    /// The sparse embedding as a `{term_id: weight}` map, empty when there is none.
110    pub fn sparse_map(&self) -> std::collections::HashMap<u32, f32> {
111        self.sparse
112            .as_ref()
113            .map(SparseVector::to_map)
114            .unwrap_or_default()
115    }
116
117    /// The multivector widened to `f32` rows, empty when there is none.
118    pub fn multivector_f32(&self) -> Vec<Vec<f32>> {
119        self.multivector
120            .as_ref()
121            .map(Multivector::to_f32)
122            .unwrap_or_default()
123    }
124}
125
126/// One scored candidate.
127#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
128pub struct ScoreEntry {
129    pub item_id: String,
130    pub score: f64,
131    /// Zero-based; `0` is the most relevant candidate.
132    pub rank: u32,
133}
134
135/// Units consumed by a score call.
136#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
137pub struct ScoreUsage {
138    pub input_tokens: u64,
139    #[serde(default, skip_serializing_if = "Option::is_none")]
140    pub images: Option<u64>,
141}
142
143/// The result of scoring one query against a candidate set.
144#[derive(Debug, Clone, Default, PartialEq)]
145pub struct ScoreResult {
146    pub model: String,
147    pub query_id: Option<String>,
148    /// Ranked candidates, best first.
149    pub scores: Vec<ScoreEntry>,
150    pub usage: Option<ScoreUsage>,
151    pub request: Option<RequestMetadata>,
152}
153
154/// A span or region the model recognised.
155#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
156pub struct Entity {
157    pub text: String,
158    pub label: String,
159    pub score: f64,
160    /// Character offset into the item text, for text models.
161    #[serde(default, skip_serializing_if = "Option::is_none")]
162    pub start: Option<i64>,
163    #[serde(default, skip_serializing_if = "Option::is_none")]
164    pub end: Option<i64>,
165    /// `[x, y, w, h]` in pixels, for visual models.
166    #[serde(default, skip_serializing_if = "Option::is_none")]
167    pub bbox: Option<Vec<i64>>,
168}
169
170/// A directed relation between two entities.
171#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
172pub struct Relation {
173    pub head: String,
174    pub tail: String,
175    pub relation: String,
176    pub score: f64,
177}
178
179/// A label with a confidence.
180#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
181pub struct Classification {
182    pub label: String,
183    pub score: f64,
184}
185
186/// A detected object with its bounding box.
187#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
188pub struct DetectedObject {
189    pub label: String,
190    pub score: f64,
191    /// `[x, y, w, h]` in pixels.
192    pub bbox: Vec<i64>,
193}
194
195/// A per-item extraction failure.
196#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
197pub struct ExtractItemError {
198    pub code: String,
199    pub message: String,
200}
201
202/// One extracted item. Which fields are populated depends on the model.
203#[derive(Debug, Clone, Default, PartialEq)]
204pub struct ExtractResult {
205    pub model: Option<String>,
206    pub id: Option<String>,
207    pub entities: Vec<Entity>,
208    pub relations: Vec<Relation>,
209    pub classifications: Vec<Classification>,
210    pub objects: Vec<DetectedObject>,
211    /// Free-form structured output, for models driven by an output schema.
212    pub data: Option<Value>,
213    /// Set when this item failed while others in the batch succeeded.
214    pub error: Option<ExtractItemError>,
215    pub request: Option<RequestMetadata>,
216}
217
218#[cfg(test)]
219mod tests {
220    // These assertions are about exact values, so exact comparison is the point.
221    #![allow(clippy::float_cmp)]
222
223    use super::*;
224
225    #[test]
226    fn sparse_converts_to_a_term_weight_map() {
227        let sparse = SparseVector {
228            indices: vec![3, 9],
229            values: vec![0.5, 0.25],
230        };
231        let map = sparse.to_map();
232        assert_eq!(map.len(), 2);
233        assert_eq!(map[&9], 0.25);
234        assert_eq!(sparse.len(), 2);
235        assert!(SparseVector::default().is_empty());
236    }
237
238    #[test]
239    fn encode_result_accessors_are_lenient_except_where_they_promise_not_to_be() {
240        let empty = EncodeResult::default();
241        assert!(empty.require_dense().is_err());
242        assert!(empty.sparse_map().is_empty());
243        assert!(empty.multivector_f32().is_empty());
244
245        let filled = EncodeResult {
246            dense: Some(vec![0.5, 0.25]),
247            sparse: Some(SparseVector {
248                indices: vec![4],
249                values: vec![1.0],
250            }),
251            ..EncodeResult::default()
252        };
253        assert_eq!(filled.require_dense().unwrap(), &[0.5, 0.25]);
254        assert_eq!(filled.sparse_map()[&4], 1.0);
255    }
256
257    #[test]
258    fn multivector_reports_its_shape_and_widens_on_request() {
259        let mv = Multivector::F16(vec![
260            vec![f16::from_f32(1.0), f16::from_f32(0.5)],
261            vec![f16::from_f32(0.0), f16::from_f32(-1.0)],
262        ]);
263        assert_eq!(mv.len(), 2);
264        assert_eq!(mv.dims(), 2);
265        assert_eq!(mv.to_f32(), vec![vec![1.0, 0.5], vec![0.0, -1.0]]);
266        assert!(Multivector::F32(Vec::new()).is_empty());
267    }
268}
sie_sdk/types/encode.rs

sie_sdk/types/
encode.rs
