aleph_alpha_client/semantic_embedding/
embedding_batch.rs

1use crate::semantic_embedding::{RequestBody, DEFAULT_EMBEDDING_MODEL};
2use crate::{Job, Prompt, SemanticRepresentation};
3use serde::{Deserialize, Serialize};
4
5const ENDPOINT: &str = "batch_semantic_embed";
6
7/// Create embeddings for multiple prompts
8#[derive(Serialize, Debug)]
9pub struct TaskBatchSemanticEmbedding<'a> {
10    /// The prompt (usually text) to be embedded.
11    pub prompts: Vec<Prompt<'a>>,
12    /// Semantic representation to embed the prompt with. This parameter is governed by the specific
13    /// use case in mind.
14    pub representation: SemanticRepresentation,
15    /// Default behaviour is to return the full embedding, but you can optionally request an
16    /// embedding compressed to a smaller set of dimensions. A size of `128` is supported for every
17    /// model.
18    ///
19    /// The 128 size is expected to have a small drop in accuracy performance (4-6%), with the
20    /// benefit of being much smaller, which makes comparing these embeddings much faster for use
21    /// cases where speed is critical.
22    ///
23    /// The 128 size can also perform better if you are embedding short texts or documents.
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub compress_to_size: Option<u32>,
26}
27
28/// Heap allocated vec of embeddings. Can hold full embeddings or compressed ones
29#[derive(Deserialize)]
30pub struct BatchSemanticEmbeddingOutput {
31    pub embeddings: Vec<Vec<f32>>,
32}
33
34impl Job for TaskBatchSemanticEmbedding<'_> {
35    type Output = BatchSemanticEmbeddingOutput;
36    type ResponseBody = BatchSemanticEmbeddingOutput;
37
38    fn build_request(&self, client: &reqwest::Client, base: &str) -> reqwest::RequestBuilder {
39        let body = RequestBody {
40            model: DEFAULT_EMBEDDING_MODEL,
41            semantic_embedding_task: self,
42        };
43        client.post(format!("{base}/{ENDPOINT}")).json(&body)
44    }
45
46    fn body_to_output(&self, response: Self::ResponseBody) -> Self::Output {
47        response
48    }
49}
aleph_alpha_client/semantic_embedding/embedding_batch.rs

aleph_alpha_client/semantic_embedding/
embedding_batch.rs
