sie_sdk/
ndarray.rs

1//! Conversions into [`ndarray`] types, for callers whose numeric code already uses it.
2//!
3//! Enabled by the `ndarray` cargo feature. The core types stay `Vec`-based so the crate
4//! does not force a numeric dependency on anyone who does not want one.
5
6use half::f16;
7use ndarray::{Array1, Array2};
8
9use crate::error::{Error, Result};
10use crate::types::{EncodeResult, Multivector, SparseVector};
11
12/// The dense embedding as a 1-D array.
13pub fn dense_array(result: &EncodeResult) -> Result<Array1<f32>> {
14    Ok(Array1::from(result.require_dense()?.to_vec()))
15}
16
17/// A multivector as a `[tokens, dims]` array of `f32`.
18pub fn multivector_f32(multivector: &Multivector) -> Result<Array2<f32>> {
19    rows_to_array(multivector.to_f32())
20}
21
22/// A multivector as a `[tokens, dims]` array of `f16`, when the wire carried `f16`.
23///
24/// Returns `None` for an `f32` multivector rather than narrowing it, which would lose
25/// precision the server chose to send.
26pub fn multivector_f16(multivector: &Multivector) -> Option<Result<Array2<f16>>> {
27    match multivector {
28        Multivector::F16(rows) => Some(rows_to_array(rows.clone())),
29        Multivector::F32(_) => None,
30    }
31}
32
33/// A sparse embedding as a dense array of `dims` elements.
34///
35/// Every index must be inside `dims`; a term id past the end means the caller's `dims` does
36/// not match the model's vocabulary.
37pub fn sparse_to_dense(sparse: &SparseVector, dims: usize) -> Result<Array1<f32>> {
38    let mut dense = Array1::zeros(dims);
39    for (index, value) in sparse.indices.iter().zip(&sparse.values) {
40        let index = *index as usize;
41        if index >= dims {
42            return Err(Error::decode(format!(
43                "sparse index {index} is outside a {dims}-dimensional vector"
44            )));
45        }
46        dense[index] = *value;
47    }
48    Ok(dense)
49}
50
51fn rows_to_array<T>(rows: Vec<Vec<T>>) -> Result<Array2<T>> {
52    let height = rows.len();
53    let width = rows.first().map_or(0, Vec::len);
54    if rows.iter().any(|row| row.len() != width) {
55        return Err(Error::decode("multivector rows have inconsistent widths"));
56    }
57    Array2::from_shape_vec((height, width), rows.into_iter().flatten().collect())
58        .map_err(|err| Error::decode(format!("could not shape the multivector: {err}")))
59}
60
61#[cfg(test)]
62mod tests {
63    // These assertions are about exact values, so exact comparison is the point.
64    #![allow(clippy::float_cmp)]
65
66    use super::*;
67
68    #[test]
69    fn dense_becomes_a_one_dimensional_array() {
70        let result = EncodeResult {
71            dense: Some(vec![0.5, -0.25, 1.0]),
72            ..EncodeResult::default()
73        };
74        let array = dense_array(&result).unwrap();
75        assert_eq!(array.len(), 3);
76        assert_eq!(array[1], -0.25);
77        assert!(dense_array(&EncodeResult::default()).is_err());
78    }
79
80    #[test]
81    fn multivectors_keep_their_shape() {
82        let multivector = Multivector::F32(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]);
83        let array = multivector_f32(&multivector).unwrap();
84        assert_eq!(array.shape(), &[2, 3]);
85        assert_eq!(array[[1, 2]], 6.0);
86    }
87
88    #[test]
89    fn f16_multivectors_are_offered_without_narrowing_f32_ones() {
90        let narrow = Multivector::F16(vec![vec![f16::from_f32(1.5), f16::from_f32(0.5)]]);
91        let array = multivector_f16(&narrow).unwrap().unwrap();
92        assert_eq!(array.shape(), &[1, 2]);
93        assert_eq!(array[[0, 0]], f16::from_f32(1.5));
94
95        assert!(multivector_f16(&Multivector::F32(vec![vec![1.0]])).is_none());
96    }
97
98    #[test]
99    fn ragged_multivectors_are_rejected_rather_than_reshaped() {
100        let ragged = Multivector::F32(vec![vec![1.0, 2.0], vec![3.0]]);
101        assert!(multivector_f32(&ragged).is_err());
102    }
103
104    #[test]
105    fn sparse_expands_into_a_dense_vector() {
106        let sparse = SparseVector {
107            indices: vec![0, 3],
108            values: vec![1.0, 0.5],
109        };
110        let dense = sparse_to_dense(&sparse, 4).unwrap();
111        assert_eq!(dense.to_vec(), vec![1.0, 0.0, 0.0, 0.5]);
112        assert!(sparse_to_dense(&sparse, 3).is_err());
113    }
114}
sie_sdk/ndarray.rs

sie_sdk/
ndarray.rs
