Skip to main content

urna_runtime/
dtype.rs

1//! Runtime view of the embeddings section dtype, carved out of
2//! `mmap_file.rs` so that file stays under the 300-line crate guard.
3
4use urna_format::UrnaError;
5
6use crate::error::RuntimeError;
7
8/// Runtime view of the embeddings section dtype.
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum DType {
11    Float32,
12    Float16,
13    Int8,
14    Int4,
15}
16
17impl DType {
18    pub(crate) fn from_str(s: &str) -> Result<Self, RuntimeError> {
19        match s {
20            "float32" => Ok(Self::Float32),
21            "float16" => Ok(Self::Float16),
22            "int8" => Ok(Self::Int8),
23            "int4" => Ok(Self::Int4),
24            other => Err(RuntimeError::Format(UrnaError::UnsupportedDType(
25                other.into(),
26            ))),
27        }
28    }
29    /// Nominal on-disk bytes per stored embedding value. int4 packs two
30    /// codes per byte (rounds to 0 here); the exact section size, with the
31    /// f16 group scales, is `expected_embeddings_size`.
32    pub fn bytes_per_value(self) -> usize {
33        match self {
34            Self::Float32 => 4,
35            Self::Float16 => 2,
36            Self::Int8 => 1,
37            Self::Int4 => 0,
38        }
39    }
40    pub fn name(self) -> &'static str {
41        match self {
42            Self::Float32 => "float32",
43            Self::Float16 => "float16",
44            Self::Int8 => "int8",
45            Self::Int4 => "int4",
46        }
47    }
48}