combs-formats 0.2.2

//! HuggingFace safetensors adapter: `config.json` + `model.safetensors`
//! (single-file or sharded with `model.safetensors.index.json`), mmap-backed.

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use memmap2::Mmap;
use safetensors::SafeTensors;

use crate::metadata::ModelMetadata;
use crate::source::{ModelSource, SamplerConfig, TensorDtype, TensorReader};
use crate::tokenizer::TokenizerSpec;
use crate::{FormatError, Result};

/// One memory-mapped safetensors shard.
struct Shard {
    mmap: Mmap,
}

/// Lightweight per-tensor index entry (populated once at load).
struct TensorEntry {
    shard: usize,
    dtype: TensorDtype,
    shape: Vec<usize>,
}

/// [`ModelSource`] over a HuggingFace-format directory.
///
/// Standard layout:
///
/// ```text
/// <dir>/config.json                   (required)
/// <dir>/generation_config.json        (optional)
/// <dir>/tokenizer.json                (required by the runtime)
/// <dir>/tokenizer_config.json         (optional; chat special tokens)
/// <dir>/model.safetensors             (single-file), or
/// <dir>/model.safetensors.index.json  (sharded)
/// ```
///
/// For diffusion sub-dirs (`unet/`, `vae/`, `text_encoder/`) use
/// [`SafetensorsSource::load_weights_only`], which tolerates
/// `diffusion_pytorch_model.safetensors` (or `model.safetensors`) and no
/// tokenizer/config.
///
/// Files are memory-mapped; [`ModelSource::open_tensor`] returns zero-copy
/// views into the mapping.
pub struct SafetensorsSource {
    metadata: ModelMetadata,
    tokenizer: TokenizerSpec,
    sampler: Option<SamplerConfig>,
    shards: Vec<Shard>,
    index: HashMap<String, TensorEntry>,
}

fn read_json(path: &Path, required: bool) -> Result<Option<serde_json::Value>> {
    match std::fs::read_to_string(path) {
        Ok(text) => serde_json::from_str(&text)
            .map(Some)
            .map_err(|source| FormatError::Json {
                context: path.display().to_string(),
                source,
            }),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound && !required => Ok(None),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
            Err(FormatError::MissingFile(path.display().to_string()))
        }
        Err(e) => Err(FormatError::Io(e)),
    }
}

fn convert_dtype(
    _name: &str,
    dtype: safetensors::Dtype,
) -> Result<TensorDtype> {
    match dtype {
        safetensors::Dtype::F64 => Ok(TensorDtype::F64),
        safetensors::Dtype::F32 => Ok(TensorDtype::F32),
        safetensors::Dtype::F16 => Ok(TensorDtype::F16),
        safetensors::Dtype::BF16 => Ok(TensorDtype::BF16),
        safetensors::Dtype::I64 => Ok(TensorDtype::I64),
        safetensors::Dtype::I32 => Ok(TensorDtype::I32),
        safetensors::Dtype::I16 => Ok(TensorDtype::I16),
        safetensors::Dtype::I8 => Ok(TensorDtype::I8),
        safetensors::Dtype::U64 => Ok(TensorDtype::U64),
        safetensors::Dtype::U32 => Ok(TensorDtype::U32),
        safetensors::Dtype::U16 => Ok(TensorDtype::U16),
        safetensors::Dtype::U8 => Ok(TensorDtype::U8),
        safetensors::Dtype::BOOL => Ok(TensorDtype::Bool),
        other => Err(FormatError::UnsupportedDtype {
            tensor: _name.to_string(),
            dtype: format!("{other:?}"),
        }),
    }
}

impl SafetensorsSource {
    /// Opens a model directory (see type docs for the expected layout).
    pub fn load(dir: impl AsRef<Path>) -> Result<Self> {
        let dir = dir.as_ref();

        // --- config + metadata -------------------------------------------------
        let config = read_json(&dir.join("config.json"), true)?.expect("required");
        let generation_config = read_json(&dir.join("generation_config.json"), false)?;
        let metadata =
            ModelMetadata::from_hf_config(&config, generation_config.as_ref())?;

        // --- tokenizer ----------------------------------------------------------
        let tokenizer_json = dir.join("tokenizer.json");
        if !tokenizer_json.exists() {
            return Err(FormatError::MissingFile(
                tokenizer_json.display().to_string(),
            ));
        }
        let mut added_tokens = HashMap::new();
        let mut chat_template = None;
        let mut add_bos = None;
        if let Some(tc) = read_json(&dir.join("tokenizer_config.json"), false)? {
            if let Some(map) = tc.get("added_tokens_decoder").and_then(|v| v.as_object()) {
                for (id, entry) in map {
                    if let (Ok(id), Some(content)) = (
                        id.parse::<u32>(),
                        entry.get("content").and_then(|c| c.as_str()),
                    ) {
                        added_tokens.insert(id, content.to_string());
                    }
                }
            }
            chat_template = tc
                .get("chat_template")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());
            add_bos = tc.get("add_bos_token").and_then(|v| v.as_bool());
        }
        let tokenizer = TokenizerSpec {
            tokenizer_json,
            added_tokens,
            chat_template,
            add_bos,
        };

        // --- sampler defaults ----------------------------------------------------
        let sampler = generation_config.as_ref().map(|gc| SamplerConfig {
            temperature: gc
                .get("temperature")
                .and_then(|v| v.as_f64())
                .map(|v| v as f32),
            top_p: gc.get("top_p").and_then(|v| v.as_f64()).map(|v| v as f32),
            top_k: gc
                .get("top_k")
                .and_then(|v| v.as_u64())
                .map(|v| v as usize),
            repetition_penalty: gc
                .get("repetition_penalty")
                .and_then(|v| v.as_f64())
                .map(|v| v as f32),
            max_new_tokens: gc
                .get("max_new_tokens")
                .or_else(|| gc.get("max_length"))
                .and_then(|v| v.as_u64())
                .map(|v| v as usize),
        });

        // --- weight shards ---------------------------------------------------------
        let shard_files = Self::collect_shard_files(dir, &["model.safetensors"])?;
        let (shards, index) = Self::mmap_shards(&shard_files)?;

        Ok(SafetensorsSource {
            metadata,
            tokenizer,
            sampler,
            shards,
            index,
        })
    }

    /// Opens a directory purely for weight loading. No `config.json` or
    /// `tokenizer.json` is required; `metadata()` returns a placeholder and
    /// `tokenizer()` errors. Accepts diffusion-style
    /// `diffusion_pytorch_model.safetensors` as well as `model.safetensors`.
    pub fn load_weights_only(dir: impl AsRef<Path>, architecture: &str) -> Result<Self> {
        let dir = dir.as_ref();
        let shard_files = Self::collect_shard_files(
            dir,
            &[
                "diffusion_pytorch_model.safetensors",
                "model.safetensors",
            ],
        )?;
        let (shards, index) = Self::mmap_shards(&shard_files)?;
        Ok(SafetensorsSource {
            metadata: ModelMetadata::diffusion_placeholder(architecture),
            tokenizer: TokenizerSpec::placeholder(),
            sampler: None,
            shards,
            index,
        })
    }

    fn collect_shard_files(dir: &Path, base_names: &[&str]) -> Result<Vec<PathBuf>> {
        for base in base_names {
            let index_json = dir.join(format!("{base}.index.json"));
            if index_json.exists() {
                let idx = read_json(&index_json, true)?.expect("required");
                let weight_map = idx
                    .get("weight_map")
                    .and_then(|v| v.as_object())
                    .ok_or_else(|| FormatError::MissingField("weight_map".to_string()))?;
                let mut files: Vec<PathBuf> = weight_map
                    .values()
                    .filter_map(|v| v.as_str())
                    .map(|f| dir.join(f))
                    .collect();
                files.sort();
                files.dedup();
                return Ok(files);
            }
            let single = dir.join(base);
            if single.exists() {
                return Ok(vec![single]);
            }
        }
        Err(FormatError::MissingFile(format!(
            "{:?} (or .index.json)",
            base_names
                .iter()
                .map(|b| dir.join(b))
                .collect::<Vec<_>>()
        )))
    }

    fn mmap_shards(shard_files: &[PathBuf]) -> Result<(Vec<Shard>, HashMap<String, TensorEntry>)> {
        let mut shards = Vec::with_capacity(shard_files.len());
        let mut index = HashMap::new();
        for (shard_idx, file) in shard_files.iter().enumerate() {
            let f = std::fs::File::open(file)?;
            // SAFETY: the file is opened read-only and never mutated by us;
            // external mutation of an mmap'd model file is out of scope.
            let mmap = unsafe { Mmap::map(&f)? };
            let st = SafeTensors::deserialize(&mmap).map_err(|e| {
                FormatError::Safetensors(format!("{}: {e}", file.display()))
            })?;
            for (name, view) in st.tensors() {
                index.insert(
                    name.clone(),
                    TensorEntry {
                        shard: shard_idx,
                        dtype: convert_dtype(&name, view.dtype())?,
                        shape: view.shape().to_vec(),
                    },
                );
            }
            shards.push(Shard { mmap });
        }
        Ok((shards, index))
    }
}

impl ModelSource for SafetensorsSource {
    fn metadata(&self) -> &ModelMetadata {
        &self.metadata
    }

    fn tensor_names(&self) -> Vec<String> {
        let mut names: Vec<String> = self.index.keys().cloned().collect();
        names.sort();
        names
    }

    fn open_tensor(&self, name: &str) -> Result<TensorReader<'_>> {
        let entry = self
            .index
            .get(name)
            .ok_or_else(|| FormatError::TensorNotFound(name.to_string()))?;
        let shard = &self.shards[entry.shard];
        let st = SafeTensors::deserialize(&shard.mmap)
            .map_err(|e| FormatError::Safetensors(e.to_string()))?;
        let view = st
            .tensor(name)
            .map_err(|e| FormatError::Safetensors(e.to_string()))?;
        Ok(TensorReader::new(
            name.to_string(),
            entry.shape.clone(),
            entry.dtype,
            view.data(),
        ))
    }

    fn tokenizer(&self) -> Result<TokenizerSpec> {
        if self.tokenizer.tokenizer_json.as_os_str().is_empty() {
            return Err(FormatError::MissingFile(
                "tokenizer.json (weights-only source has no tokenizer)".to_string(),
            ));
        }
        Ok(self.tokenizer.clone())
    }

    fn sampler_defaults(&self) -> Option<SamplerConfig> {
        self.sampler.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Builds a tiny synthetic HF model dir and checks the adapter surface.
    #[test]
    fn lists_and_reads_tensors() {
        let dir = tempfile::tempdir().unwrap();
        let root = dir.path();

        std::fs::write(
            root.join("config.json"),
            serde_json::to_string(&serde_json::json!({
                "model_type": "llama",
                "hidden_size": 8,
                "intermediate_size": 16,
                "num_hidden_layers": 1,
                "num_attention_heads": 2,
                "num_key_value_heads": 1,
                "vocab_size": 32,
                "max_position_embeddings": 128,
                "rope_theta": 10000,
                "rms_norm_eps": 1e-5,
                "tie_word_embeddings": true,
                "eos_token_id": 0,
                "bos_token_id": 0
            }))
            .unwrap(),
        )
        .unwrap();
        std::fs::write(root.join("tokenizer.json"), "{}").unwrap();
        std::fs::write(
            root.join("tokenizer_config.json"),
            serde_json::to_string(&serde_json::json!({
                "added_tokens_decoder": {
                    "0": {"content": "<|endoftext|>", "special": true},
                    "2": {"content": "<|im_end|>", "special": true}
                }
            }))
            .unwrap(),
        )
        .unwrap();

        // Two tensors: one F32, one BF16.
        let w1_bytes: Vec<u8> = (0..16)
            .flat_map(|i| (i as f32).to_le_bytes())
            .collect();
        let w2_bytes: Vec<u8> = (0..8)
            .flat_map(|i| half::bf16::from_f32(i as f32 * 0.5).to_le_bytes())
            .collect();
        let v1 = safetensors::tensor::TensorView::new(
            safetensors::Dtype::F32,
            vec![4, 4],
            &w1_bytes,
        )
        .unwrap();
        let v2 = safetensors::tensor::TensorView::new(
            safetensors::Dtype::BF16,
            vec![8],
            &w2_bytes,
        )
        .unwrap();

        // Add an I64 position_ids buffer like some CLIP text encoders ship.
        let pos_ids_bytes: Vec<u8> = (0..8i64)
            .flat_map(|i| i.to_le_bytes())
            .collect();
        let v_pos = safetensors::tensor::TensorView::new(
            safetensors::Dtype::I64,
            vec![8],
            &pos_ids_bytes,
        )
        .unwrap();

        safetensors::serialize_to_file(
            vec![("a.weight", v1), ("b.weight", v2), ("text_model.embeddings.position_ids", v_pos)],
            None,
            root.join("model.safetensors").as_path(),
        )
        .unwrap();

        let src = SafetensorsSource::load(root).unwrap();
        assert_eq!(src.metadata().architecture, "llama");
        assert_eq!(src.metadata().head_dim, 4);
        assert_eq!(
            src.tensor_names(),
            vec!["a.weight", "b.weight", "text_model.embeddings.position_ids"]
        );

        let pos = src.open_tensor("text_model.embeddings.position_ids").unwrap();
        assert_eq!(pos.dtype(), TensorDtype::I64);
        let pos_vals: Vec<f32> = pos.load_data().unwrap().to_vec().unwrap();
        assert_eq!(pos_vals[7], 7.0);

        let a = src.open_tensor("a.weight").unwrap();
        assert_eq!(a.shape(), &[4, 4]);
        assert_eq!(a.dtype(), TensorDtype::F32);
        let data = a.load_data().unwrap();
        let vals: Vec<f32> = data.to_vec().unwrap();
        assert_eq!(vals[15], 15.0);

        let b = src.open_tensor("b.weight").unwrap();
        assert_eq!(b.dtype(), TensorDtype::BF16);
        let vals: Vec<f32> = b.load_data().unwrap().to_vec().unwrap();
        assert!((vals[3] - 1.5).abs() < 1e-3);

        let tok = src.tokenizer().unwrap();
        assert_eq!(tok.special_token_id("<|im_end|>"), Some(2));

        assert!(src.open_tensor("missing").is_err());
    }

    /// `load_weights_only` accepts diffusion-style filenames and no config.
    #[test]
    fn loads_weights_only_without_config_or_tokenizer() {
        let dir = tempfile::tempdir().unwrap();
        let root = dir.path();

        let bytes: Vec<u8> = (0..16)
            .flat_map(|i| (i as f32).to_le_bytes())
            .collect();
        let t = safetensors::tensor::TensorView::new(
            safetensors::Dtype::F32,
            vec![4, 4],
            &bytes,
        )
        .unwrap();
        safetensors::serialize_to_file(
            vec![("conv_in.weight", t)],
            None,
            root.join("diffusion_pytorch_model.safetensors")
                .as_path(),
        )
        .unwrap();

        let src =
            SafetensorsSource::load_weights_only(root, "stable-diffusion-unet").unwrap();
        assert_eq!(src.metadata().architecture, "stable-diffusion-unet");
        assert_eq!(src.tensor_names(), vec!["conv_in.weight"]);
        assert!(src.tokenizer().is_err());
    }
}