ik_llama_cpp_2/context/
embeddings.rs

1//! Embeddings + reranker support: the `llama_get_embeddings*` accessors.
2//!
3//! Embeddings are produced when the context is built with
4//! `LlamaContextParams::with_embeddings(true)` and (usually) a pooling type set
5//! via `with_pooling_type`. Fetch a per-token vector with
6//! [`LlamaContext::embeddings_ith`] or a pooled per-sequence vector with
7//! [`LlamaContext::embeddings_seq_ith`]; each slice is `n_embd` long.
8//!
9//! A **reranker** is just an embedding model run with a rank/classification
10//! pooling head: the pooled "embedding" of a query+document sequence is a
11//! relevance score. Read it with [`LlamaContext::embeddings_seq_ith`] after a
12//! decode (with a rank pooling type, the slice collapses to the score). ik does
13//! not expose a dedicated `LLAMA_POOLING_TYPE_RANK` constant — the available
14//! pooling types are `NONE`, `MEAN`, `CLS`, and `LAST`.
15
16use ik_llama_cpp_sys as sys;
17
18use crate::context::LlamaContext;
19
20/// Errors returned by the embedding accessors.
21#[derive(Debug, PartialEq, Eq, thiserror::Error)]
22pub enum EmbeddingsError {
23    /// The context was created without `with_embeddings(true)`.
24    #[error("embeddings were not enabled in the context params")]
25    NotEnabled,
26    /// No embeddings for the given token: its batch entry did not request
27    /// logits/embeddings, or the index was out of range.
28    #[error("no embeddings available for token index {0}")]
29    LogitsNotEnabled(i32),
30    /// No per-sequence embeddings for the given sequence: the pooling type is
31    /// `LLAMA_POOLING_TYPE_NONE`, or the sequence id was out of range.
32    #[error("no sequence embeddings for seq {0} (pooling type NONE, or out of range)")]
33    NonePoolType(i32),
34}
35
36impl LlamaContext<'_> {
37    /// The model's embedding dimension (`n_embd`).
38    ///
39    /// ik exposes only `llama_model_n_embd` (there is no plain `llama_n_embd`),
40    /// so we resolve the model from the context pointer first. The context does
41    /// not hold a model reference, hence the round-trip through `llama_get_model`.
42    #[must_use]
43    fn n_embd(&self) -> usize {
44        // SAFETY: the context owns a valid model pointer for its whole lifetime.
45        let n = unsafe { sys::llama_model_n_embd(sys::llama_get_model(self.context.as_ptr())) };
46        usize::try_from(n).unwrap_or(0)
47    }
48
49    /// Get the embeddings for the `i`th token in the current context.
50    ///
51    /// # Returns
52    ///
53    /// A slice with the embeddings for the last decoded batch of the given
54    /// token. The size corresponds to the `n_embd` of the context's model.
55    ///
56    /// # Errors
57    ///
58    /// - [`EmbeddingsError::NotEnabled`] if the context was built without
59    ///   `with_embeddings(true)`.
60    /// - [`EmbeddingsError::LogitsNotEnabled`] if that token was not decoded
61    ///   with logits enabled (ik returns a null pointer), or `i` is out of range.
62    pub fn embeddings_ith(&self, i: i32) -> Result<&[f32], EmbeddingsError> {
63        if !self.raw_params.embeddings {
64            return Err(EmbeddingsError::NotEnabled);
65        }
66        let n_embd = self.n_embd();
67        // SAFETY: ptr valid until the next decode; slice length is n_embd.
68        let ptr = unsafe { sys::llama_get_embeddings_ith(self.context.as_ptr(), i) };
69        if ptr.is_null() {
70            return Err(EmbeddingsError::LogitsNotEnabled(i));
71        }
72        Ok(unsafe { std::slice::from_raw_parts(ptr, n_embd) })
73    }
74
75    /// Get the pooled embeddings for the `seq`th sequence in the current context.
76    ///
77    /// # Returns
78    ///
79    /// A slice with the pooled embeddings for the last decoded batch. The size
80    /// corresponds to the `n_embd` of the context's model.
81    ///
82    /// # Errors
83    ///
84    /// - [`EmbeddingsError::NotEnabled`] if the context was built without
85    ///   `with_embeddings(true)`.
86    /// - [`EmbeddingsError::NonePoolType`] if the model uses
87    ///   `LLAMA_POOLING_TYPE_NONE` (ik returns a null pointer), or `seq` exceeds
88    ///   the max sequence id.
89    pub fn embeddings_seq_ith(&self, seq: i32) -> Result<&[f32], EmbeddingsError> {
90        if !self.raw_params.embeddings {
91            return Err(EmbeddingsError::NotEnabled);
92        }
93        let n_embd = self.n_embd();
94        // SAFETY: ptr valid until the next decode; slice length is n_embd.
95        let ptr = unsafe { sys::llama_get_embeddings_seq(self.context.as_ptr(), seq) };
96        if ptr.is_null() {
97            return Err(EmbeddingsError::NonePoolType(seq));
98        }
99        Ok(unsafe { std::slice::from_raw_parts(ptr, n_embd) })
100    }
101}
ik_llama_cpp_2/context/embeddings.rs

ik_llama_cpp_2/context/
embeddings.rs
