Skip to main content

corpus_ipc/models/
embeddings.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use serde::{Deserialize, Serialize};
4
5use crate::validation::{
6    ProtocolLimits, Validate, ValidationError, add_to_total, check_count, check_finite_slice,
7    check_opt_string,
8};
9
10use super::de::{de_embedding, de_opt_session_id};
11
12/// Batch of embeddings for projector and transformer components.
13#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)]
14#[serde(try_from = "EmbeddingBatchWire")]
15pub struct EmbeddingBatch {
16    /// Optional session ID for concurrent experiment isolation.
17    pub session_id: Option<String>,
18    /// Unique batch identifier for correlation.
19    pub batch_id: u64,
20    /// Embedding vector from compute processing.
21    pub embedding: Vec<f32>,
22    /// Sequence length for transformer compatibility.
23    pub sequence_length: usize,
24}
25
26#[derive(Deserialize)]
27struct EmbeddingBatchWire {
28    #[serde(deserialize_with = "de_opt_session_id")]
29    session_id: Option<String>,
30    batch_id: u64,
31    #[serde(deserialize_with = "de_embedding")]
32    embedding: Vec<f32>,
33    sequence_length: usize,
34}
35
36impl TryFrom<EmbeddingBatchWire> for EmbeddingBatch {
37    type Error = ValidationError;
38
39    fn try_from(wire: EmbeddingBatchWire) -> Result<Self, Self::Error> {
40        let batch = EmbeddingBatch {
41            session_id: wire.session_id,
42            batch_id: wire.batch_id,
43            embedding: wire.embedding,
44            sequence_length: wire.sequence_length,
45        };
46        batch.validate()?;
47        Ok(batch)
48    }
49}
50
51impl Validate for EmbeddingBatch {
52    fn validate_with(&self, limits: ProtocolLimits) -> Result<(), ValidationError> {
53        check_opt_string("session_id", self.session_id.as_deref(), limits)?;
54        check_count("embedding", self.embedding.len(), limits.max_channel_values)?;
55        check_count(
56            "sequence_length",
57            self.sequence_length,
58            limits.max_aggregate_records,
59        )?;
60        check_finite_slice("embedding", &self.embedding)?;
61        let mut total = 0;
62        add_to_total(&mut total, self.embedding.len(), limits, "aggregate")?;
63        Ok(())
64    }
65}