Skip to main content

hf_fetch_model/
inspect.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Safetensors header inspection (local and remote).
4//!
5//! Reads tensor metadata (names, shapes, dtypes, byte offsets) from
6//! `.safetensors` files without downloading full weight data. Supports
7//! cache-first resolution with HTTP Range request fallback.
8//!
9//! The primary types are [`TensorInfo`] (per-tensor metadata),
10//! [`SafetensorsHeaderInfo`] (parsed header), and [`ShardedIndex`]
11//! (shard-to-tensor mapping for sharded models).
12
13use std::collections::HashMap;
14use std::path::{Path, PathBuf};
15
16use serde::Serialize;
17use tokio::task::JoinSet;
18
19use crate::cache;
20use crate::cache_layout;
21use crate::chunked;
22use crate::error::FetchError;
23
24// -----------------------------------------------------------------------
25// Types
26// -----------------------------------------------------------------------
27
28/// Metadata for a single tensor from a `.safetensors` header.
29///
30/// This is hf-fetch-model's own type — lightweight, no quantization logic.
31/// Consumers (e.g., anamnesis) map this into their own richer types.
32#[derive(Debug, Clone, Serialize)]
33pub struct TensorInfo {
34    /// Tensor name (e.g., `"model.layers.0.self_attn.q_proj.weight"`).
35    pub name: String,
36    /// Element dtype string as it appears in the header (e.g., `"F8_E4M3"`, `"BF16"`).
37    pub dtype: String,
38    /// Tensor shape (e.g., `[7168, 7168]`).
39    pub shape: Vec<usize>,
40    /// Byte offset range `[start, end)` within the data section of the file.
41    pub data_offsets: (u64, u64),
42}
43
44impl TensorInfo {
45    /// Total number of elements (product of shape dimensions).
46    ///
47    /// Returns `1` for a scalar (empty shape).
48    #[must_use]
49    pub fn num_elements(&self) -> u64 {
50        self.shape.iter().fold(1u64, |acc, &d| {
51            // CAST: usize → u64, dimension values fit in u64
52            #[allow(clippy::as_conversions)]
53            let dim = d as u64;
54            acc.saturating_mul(dim)
55        })
56    }
57
58    /// Byte length of the tensor data (`end - start`).
59    #[must_use]
60    pub const fn byte_len(&self) -> u64 {
61        self.data_offsets.1.saturating_sub(self.data_offsets.0)
62    }
63
64    /// Bytes per element for the tensor's dtype, if recognized.
65    ///
66    /// Returns `None` for unknown dtype strings. Recognized dtypes:
67    ///
68    /// | Dtype string | Bytes | Notes |
69    /// |-------------|-------|-------|
70    /// | `"BOOL"` | 1 | |
71    /// | `"U8"`, `"I8"` | 1 | |
72    /// | `"F8_E4M3"`, `"F8_E5M2"` | 1 | FP8 variants |
73    /// | `"U16"`, `"I16"`, `"F16"`, `"BF16"` | 2 | |
74    /// | `"U32"`, `"I32"`, `"F32"` | 4 | |
75    /// | `"U64"`, `"I64"`, `"F64"` | 8 | |
76    #[must_use]
77    pub fn dtype_bytes(&self) -> Option<usize> {
78        // BORROW: explicit .as_str() instead of Deref coercion
79        match self.dtype.as_str() {
80            "BOOL" | "U8" | "I8" | "F8_E4M3" | "F8_E5M2" => Some(1),
81            "U16" | "I16" | "F16" | "BF16" => Some(2),
82            "U32" | "I32" | "F32" => Some(4),
83            "U64" | "I64" | "F64" => Some(8),
84            _ => None,
85        }
86    }
87}
88
89/// Quantization scheme + size estimates for a cached `.safetensors` file.
90///
91/// Populated by [`inspect_safetensors_local`] (the cache-hit path) via
92/// `anamnesis::InspectInfo::from(&header)`. Absent (`None`) when:
93/// - the safetensors file has no detected quantization (`QuantScheme::Unquantized`),
94/// - the inspect ran over HTTP Range (the remote path's bespoke parser
95///   doesn't go through anamnesis until v0.11.1), or
96/// - the file format isn't safetensors (`GGUF` / `NPZ` / `PTH` carry no
97///   quant-method metadata).
98///
99/// Decoupled from `anamnesis::QuantScheme` (a `#[non_exhaustive]` enum) so
100/// downstream library consumers (`candle-mi`, `anamnesis`) aren't forced to
101/// match every variant. The `scheme` field stores `QuantScheme`'s `Display`
102/// output (`"FineGrainedFp8"`, `"Bnb4"`, `"Gptq"`, `"Awq"`, …); consumers
103/// that need to match exact variants should call
104/// `anamnesis::parse_safetensors_header` themselves.
105#[derive(Debug, Clone, Serialize)]
106pub struct QuantInfo {
107    /// Detected quantization scheme as the `Display` form of
108    /// `anamnesis::QuantScheme` (e.g. `"FineGrainedFp8"`, `"Bnb4"`).
109    pub scheme: String,
110    /// Bytes stored on disk for tensor data (header excluded).
111    pub stored_bytes: u64,
112    /// Estimated bytes after dequantising to `BF16`. For `BnB-NF4`/`FP4`
113    /// (`U8`-packed nibbles), this is `stored_bytes × 4`; for `FP8` / `GPTQ` /
114    /// `AWQ` / `BnB-INT8` it's `num_elements × 2` summed over weight
115    /// tensors, plus passthrough tensors copied as-is. The formula lives
116    /// in `anamnesis::InspectInfo::from(&SafetensorsHeader)` — hf-fm just
117    /// reads the result.
118    pub dequantized_bytes: u64,
119}
120
121/// Parsed safetensors header metadata.
122///
123/// Marked `#[non_exhaustive]` (since v0.10.3) — the struct has been
124/// growing through v0.10.x (the `quant_info` field landed in Phase C)
125/// and will keep growing in v0.11.x. External library consumers should
126/// pattern-match with `..` or use field reads, not exhaustive struct
127/// literals.
128#[derive(Debug, Clone, Serialize)]
129#[non_exhaustive]
130pub struct SafetensorsHeaderInfo {
131    /// All tensors in the header, in the order they appear in the JSON.
132    pub tensors: Vec<TensorInfo>,
133    /// Raw `__metadata__` entries, if present.
134    ///
135    /// For quantized models, this typically contains entries like
136    /// `quant_method`, `bits`, `group_size` that consumers like anamnesis
137    /// use to distinguish GPTQ from AWQ without downloading weights.
138    pub metadata: Option<HashMap<String, String>>,
139    /// Size of the JSON header in bytes.
140    pub header_size: u64,
141    /// Total file size in bytes (header + data), if known.
142    ///
143    /// **Source:** for local files, from `std::fs::metadata().len()`. For HTTP
144    /// Range requests, extracted from the `Content-Range` response header of
145    /// the first request (`bytes 0-7/TOTAL` → `TOTAL`). This is free — no
146    /// extra request needed.
147    pub file_size: Option<u64>,
148    /// Quantization scheme + size estimates (cached safetensors only).
149    ///
150    /// Populated by [`inspect_safetensors_local`] when a non-`Unquantized`
151    /// `QuantScheme` is detected. `None` for unquantized safetensors,
152    /// `GGUF` / `NPZ` / `PTH` files, and the remote safetensors path
153    /// (until v0.11.1 migrates that path to anamnesis).
154    #[serde(skip_serializing_if = "Option::is_none")]
155    pub quant_info: Option<QuantInfo>,
156}
157
158impl SafetensorsHeaderInfo {
159    /// Total parameter count across all tensors.
160    #[must_use]
161    pub fn total_params(&self) -> u64 {
162        self.tensors
163            .iter()
164            .map(TensorInfo::num_elements)
165            .fold(0u64, u64::saturating_add)
166    }
167
168    /// Returns tensors matching a dtype string (e.g., `"F8_E4M3"`).
169    #[must_use]
170    pub fn tensors_with_dtype(&self, dtype: &str) -> Vec<&TensorInfo> {
171        self.tensors
172            .iter()
173            // BORROW: explicit .as_str() instead of Deref coercion
174            .filter(|t| t.dtype.as_str() == dtype)
175            .collect()
176    }
177
178    /// Constructs a new [`SafetensorsHeaderInfo`] from its core fields.
179    ///
180    /// Since v0.10.3 the struct is `#[non_exhaustive]` — this constructor is
181    /// the canonical way to build one from outside the `hf-fetch-model` lib
182    /// crate (e.g. the `hf-fm` binary crate, downstream consumers like
183    /// `candle-mi`). Inside the lib crate, struct-literal syntax stays
184    /// available for the inspect entry points.
185    ///
186    /// `quant_info` is typically `None`; populated only by
187    /// [`inspect_safetensors_local`] for cached, quantized safetensors files.
188    #[must_use]
189    pub fn new(
190        tensors: Vec<TensorInfo>,
191        metadata: Option<HashMap<String, String>>,
192        header_size: u64,
193        file_size: Option<u64>,
194        quant_info: Option<QuantInfo>,
195    ) -> Self {
196        Self {
197            tensors,
198            metadata,
199            header_size,
200            file_size,
201            quant_info,
202        }
203    }
204}
205
206/// The source from which a header was read.
207#[derive(Debug, Clone, Copy, PartialEq, Eq)]
208#[non_exhaustive]
209pub enum InspectSource {
210    /// Read from local cache (no network).
211    Cached,
212    /// Fetched via HTTP Range requests.
213    Remote,
214}
215
216/// Parsed `model.safetensors.index.json` for a sharded model.
217#[derive(Debug, Clone, Serialize)]
218pub struct ShardedIndex {
219    /// Mapping from tensor name to shard filename.
220    pub weight_map: HashMap<String, String>,
221    /// Ordered list of unique shard filenames.
222    pub shards: Vec<String>,
223    /// Raw metadata from the index, if present.
224    pub metadata: Option<HashMap<String, serde_json::Value>>,
225}
226
227/// `PEFT` adapter configuration parsed from `adapter_config.json`.
228///
229/// Contains the key fields that identify an adapter: the `PEFT` type,
230/// base model, `LoRA` rank and scaling parameters, and target modules.
231/// All fields are optional because adapter configs vary across `PEFT` methods.
232#[derive(Debug, Clone, Serialize)]
233pub struct AdapterConfig {
234    /// `PEFT` method type (e.g., `"LORA"`, `"ADALORA"`, `"IA3"`).
235    pub peft_type: Option<String>,
236    /// The base model this adapter was trained on.
237    pub base_model_name_or_path: Option<String>,
238    /// `LoRA` rank (the `r` parameter). Only meaningful for `LoRA`-family methods.
239    pub r: Option<u32>,
240    /// `LoRA` alpha scaling factor. Only meaningful for `LoRA`-family methods.
241    pub lora_alpha: Option<f64>,
242    /// List of model modules targeted by the adapter.
243    pub target_modules: Vec<String>,
244    /// Task type the adapter was trained for (e.g., `"CAUSAL_LM"`).
245    pub task_type: Option<String>,
246}
247
248// -----------------------------------------------------------------------
249// JSON parsing
250// -----------------------------------------------------------------------
251
252/// Raw tensor entry as it appears in the safetensors JSON header.
253#[derive(serde::Deserialize)]
254struct RawTensorEntry {
255    dtype: String,
256    shape: Vec<usize>,
257    data_offsets: (u64, u64),
258}
259
260/// Parsed tensor list and optional metadata from a safetensors header.
261type ParsedHeader = (Vec<TensorInfo>, Option<HashMap<String, String>>);
262
263/// Parses the safetensors JSON header bytes into tensor metadata.
264///
265/// Extracts the `__metadata__` key separately (if present).
266fn parse_header_json(json_bytes: &[u8], filename: &str) -> Result<ParsedHeader, FetchError> {
267    let raw: HashMap<String, serde_json::Value> =
268        serde_json::from_slice(json_bytes).map_err(|e| FetchError::SafetensorsHeader {
269            filename: filename.to_owned(),
270            reason: format!("failed to parse header JSON: {e}"),
271        })?;
272
273    let mut metadata: Option<HashMap<String, String>> = None;
274    let mut tensors = Vec::new();
275
276    for (key, value) in raw {
277        if key == "__metadata__" {
278            if let serde_json::Value::Object(obj) = value {
279                let mut meta_map = HashMap::new();
280                for (mk, mv) in obj {
281                    // BORROW: explicit .to_owned()/.to_string() for Value → String conversion
282                    let v_str = if let Some(s) = mv.as_str() {
283                        s.to_owned()
284                    } else {
285                        mv.to_string()
286                    };
287                    meta_map.insert(mk, v_str);
288                }
289                metadata = Some(meta_map);
290            }
291            continue;
292        }
293
294        let entry: RawTensorEntry =
295            serde_json::from_value(value).map_err(|e| FetchError::SafetensorsHeader {
296                filename: filename.to_owned(),
297                reason: format!("failed to parse tensor \"{key}\": {e}"),
298            })?;
299
300        tensors.push(TensorInfo {
301            name: key,
302            dtype: entry.dtype,
303            shape: entry.shape,
304            data_offsets: entry.data_offsets,
305        });
306    }
307
308    // Sort by data offset start to preserve file order.
309    tensors.sort_by_key(|t| t.data_offsets.0);
310
311    Ok((tensors, metadata))
312}
313
314// -----------------------------------------------------------------------
315// Cache resolution
316// -----------------------------------------------------------------------
317
318/// Resolves a cached file path for a given repo, revision, and filename.
319///
320/// Returns `None` if the file is not in the local cache.
321fn resolve_cached_path(repo_id: &str, revision: &str, filename: &str) -> Option<PathBuf> {
322    let cache_dir = cache::hf_cache_dir().ok()?;
323    let repo_dir = cache_layout::repo_dir(&cache_dir, repo_id);
324    let commit_hash = cache::read_ref(&repo_dir, revision)?;
325    let cached_path = cache_layout::pointer_path(&repo_dir, &commit_hash, filename);
326    if cached_path.exists() {
327        Some(cached_path)
328    } else {
329        None
330    }
331}
332
333// -----------------------------------------------------------------------
334// Local file reading
335// -----------------------------------------------------------------------
336
337/// Inspects a single `.safetensors` file's header from a local file path.
338///
339/// Reads the first `8 + header_size` bytes from disk. Does not read tensor data.
340///
341/// # Blocking I/O
342///
343/// This function performs synchronous filesystem I/O. In async contexts, wrap
344/// it in [`tokio::task::spawn_blocking`] so the calling task does not stall
345/// the runtime — particularly important on network-mounted caches (NFS/CIFS)
346/// where `read`/`stat` calls can take tens of milliseconds each.
347///
348/// # Errors
349///
350/// Returns [`FetchError::Io`] if the file cannot be read.
351/// Returns [`FetchError::SafetensorsHeader`] if the header is malformed.
352pub fn inspect_safetensors_local(path: &Path) -> Result<SafetensorsHeaderInfo, FetchError> {
353    let file_size = std::fs::metadata(path)
354        .map_err(|e| FetchError::Io {
355            path: path.to_path_buf(),
356            source: e,
357        })?
358        .len();
359
360    // BORROW: explicit .to_string_lossy() for Path → str conversion
361    let filename = path.file_name().map_or_else(
362        || path.display().to_string(),
363        |n| n.to_string_lossy().to_string(),
364    );
365
366    let file = std::fs::File::open(path).map_err(|e| FetchError::Io {
367        path: path.to_path_buf(),
368        source: e,
369    })?;
370
371    // Cache-hit path delegates to anamnesis (v0.10.3 Phase B commit 4):
372    // single source of truth for safetensors layout. The reader variant
373    // reads the 8-byte u64 prefix + the JSON header bytes from the `Read`
374    // impl itself, *without* requiring the data section to be present
375    // (it bypasses `safetensors::SafeTensors::read_metadata` for exactly
376    // this reason). Anamnesis caps the declared header length at 100 MiB
377    // internally so the worst-case allocation is bounded.
378    //
379    // The remote path's `fetch_header_bytes` + `parse_header_json` chain
380    // stays bespoke until v0.11.1, when anamnesis's reader variant adapts
381    // to an HTTP Range source.
382    let header = anamnesis::parse_safetensors_header_from_reader(file).map_err(|e| {
383        FetchError::SafetensorsHeader {
384            // BORROW: explicit .clone() for the error variant's owned String field
385            filename: filename.clone(),
386            reason: format!("failed to parse safetensors header: {e}"),
387        }
388    })?;
389    // CAST: usize → u64, anamnesis caps header_size at 100 MiB so it always fits in u64
390    #[allow(clippy::as_conversions)]
391    let header_size = header.header_size as u64;
392
393    // Phase C: derive `quant_info` before the `header.tensors` move below.
394    // `anamnesis::InspectInfo::from(&header)` iterates the already-parsed
395    // tensors and aggregates per-role byte sums — pure computation, no I/O.
396    // Unquantized models produce no `quant_info` so the renderer suppresses
397    // the new `Format:` / `Size:` lines (absence communicates full precision).
398    let quant_info = if header.scheme == anamnesis::QuantScheme::Unquantized {
399        None
400    } else {
401        let info = anamnesis::InspectInfo::from(&header);
402        Some(QuantInfo {
403            // BORROW: explicit .to_string() — anamnesis `QuantScheme` → owned `String`
404            scheme: info.format.to_string(),
405            stored_bytes: info.current_size,
406            dequantized_bytes: info.dequantized_size,
407        })
408    };
409
410    let mut tensors: Vec<TensorInfo> = header
411        .tensors
412        .into_iter()
413        .map(|t| {
414            // CAST: usize → u64, header data_offsets fit in u64 by definition (file size is u64)
415            #[allow(clippy::as_conversions)]
416            let start = t.data_offsets.0 as u64;
417            // CAST: usize → u64, same rationale as above
418            #[allow(clippy::as_conversions)]
419            let end = t.data_offsets.1 as u64;
420            TensorInfo {
421                name: t.name,
422                // BORROW: explicit .to_string() — anamnesis `Dtype` enum → owned `String`
423                dtype: t.dtype.to_string(),
424                shape: t.shape,
425                data_offsets: (start, end),
426            }
427        })
428        .collect();
429
430    // Preserve hf-fm's v0.10.2 sort order. Anamnesis returns tensors sorted
431    // alphabetically by name; hf-fm's inspect table has always been
432    // file-ordered (sorted by start offset) so users can spot first/last
433    // tensors per shard at a glance.
434    tensors.sort_by_key(|t| t.data_offsets.0);
435
436    Ok(SafetensorsHeaderInfo {
437        tensors,
438        metadata: header.metadata,
439        header_size,
440        file_size: Some(file_size),
441        quant_info,
442    })
443}
444
445// -----------------------------------------------------------------------
446// Remote fetching (HTTP Range requests)
447// -----------------------------------------------------------------------
448
449/// Fetches safetensors header bytes via two HTTP Range requests.
450///
451/// 1. `Range: bytes=0-7` → 8-byte header length (little-endian `u64`)
452/// 2. `Range: bytes=8-{8+length-1}` → JSON header
453///
454/// Returns `(json_bytes, total_file_size)`. The file size is extracted from
455/// the `Content-Range` header of the first request.
456async fn fetch_header_bytes(
457    client: &reqwest::Client,
458    url: &str,
459    filename: &str,
460) -> Result<(Vec<u8>, Option<u64>), FetchError> {
461    // Request 1: 8-byte length prefix.
462    let resp1 = client
463        .get(url)
464        .header(reqwest::header::RANGE, "bytes=0-7")
465        .send()
466        .await
467        .map_err(|e| {
468            FetchError::Http(format!("failed to fetch header length for {filename}: {e}"))
469        })?;
470
471    if !resp1.status().is_success() && resp1.status() != reqwest::StatusCode::PARTIAL_CONTENT {
472        return Err(FetchError::Http(format!(
473            "Range request for {filename} returned status {}",
474            resp1.status()
475        )));
476    }
477
478    // Extract total file size from Content-Range: bytes 0-7/{total}
479    let file_size = resp1
480        .headers()
481        .get(reqwest::header::CONTENT_RANGE)
482        .and_then(|v| v.to_str().ok())
483        .and_then(|s| s.split('/').next_back())
484        .and_then(|s| s.parse::<u64>().ok());
485
486    let len_bytes = resp1.bytes().await.map_err(|e| {
487        FetchError::Http(format!("failed to read header length for {filename}: {e}"))
488    })?;
489
490    if len_bytes.len() < 8 {
491        return Err(FetchError::SafetensorsHeader {
492            filename: filename.to_owned(),
493            reason: format!(
494                "expected 8 bytes for length prefix, got {}",
495                len_bytes.len()
496            ),
497        });
498    }
499
500    // INDEX: first 8 bytes guaranteed by length check above
501    #[allow(clippy::indexing_slicing)]
502    let header_size = u64::from_le_bytes([
503        len_bytes[0],
504        len_bytes[1],
505        len_bytes[2],
506        len_bytes[3],
507        len_bytes[4],
508        len_bytes[5],
509        len_bytes[6],
510        len_bytes[7],
511    ]);
512
513    // Request 2: JSON header.
514    let range_end = 8u64.saturating_add(header_size).saturating_sub(1);
515    let range_header = format!("bytes=8-{range_end}");
516    let resp2 = client
517        .get(url)
518        // BORROW: explicit .as_str() instead of Deref coercion
519        .header(reqwest::header::RANGE, range_header.as_str())
520        .send()
521        .await
522        .map_err(|e| {
523            FetchError::Http(format!("failed to fetch header JSON for {filename}: {e}"))
524        })?;
525
526    if !resp2.status().is_success() && resp2.status() != reqwest::StatusCode::PARTIAL_CONTENT {
527        return Err(FetchError::Http(format!(
528            "Range request for {filename} header JSON returned status {}",
529            resp2.status()
530        )));
531    }
532
533    let json_bytes = resp2
534        .bytes()
535        .await
536        .map_err(|e| FetchError::Http(format!("failed to read header JSON for {filename}: {e}")))?;
537
538    Ok((json_bytes.to_vec(), file_size))
539}
540
541// -----------------------------------------------------------------------
542// Public API: single-file inspection
543// -----------------------------------------------------------------------
544
545/// Inspects a single `.safetensors` file's header (cache-first).
546///
547/// Checks the local HF cache first. If the file is cached, reads the header
548/// from disk with zero network requests. Otherwise, falls back to two HTTP
549/// Range requests (8-byte length prefix + JSON header). Does not download
550/// tensor data in either case.
551///
552/// # Errors
553///
554/// Returns [`FetchError::Http`] if the Range requests fail.
555/// Returns [`FetchError::SafetensorsHeader`] if the header is malformed.
556pub async fn inspect_safetensors(
557    repo_id: &str,
558    filename: &str,
559    token: Option<&str>,
560    revision: Option<&str>,
561) -> Result<(SafetensorsHeaderInfo, InspectSource), FetchError> {
562    let rev = revision.unwrap_or("main");
563
564    // Try local cache first.
565    if let Some(cached_path) = resolve_cached_path(repo_id, rev, filename) {
566        let info = inspect_safetensors_local(&cached_path)?;
567        return Ok((info, InspectSource::Cached));
568    }
569
570    // Fall back to HTTP Range requests.
571    let client = chunked::build_client(token)?;
572    let url = chunked::build_download_url(repo_id, rev, filename);
573
574    // BORROW: explicit .as_str() instead of Deref coercion
575    let (json_bytes, file_size) = fetch_header_bytes(&client, url.as_str(), filename).await?;
576
577    // CAST: usize → u64, JSON buffer length is always small
578    #[allow(clippy::as_conversions)]
579    let header_size = json_bytes.len() as u64;
580
581    let (tensors, metadata) = parse_header_json(&json_bytes, filename)?;
582
583    Ok((
584        SafetensorsHeaderInfo {
585            tensors,
586            metadata,
587            header_size,
588            file_size,
589            // Remote path doesn't go through anamnesis yet; quant detection
590            // arrives in v0.11.1 when this path migrates to
591            // `anamnesis::parse_safetensors_header_from_reader` over an
592            // HTTP Range adapter.
593            quant_info: None,
594        },
595        InspectSource::Remote,
596    ))
597}
598
599/// Inspects a single `.safetensors` file from cache only.
600///
601/// Resolves the file in the local HF cache using the given `repo_id`,
602/// `revision`, and `filename`. Returns an error if the file is not cached.
603///
604/// # Blocking I/O
605///
606/// Performs synchronous filesystem I/O; wrap in [`tokio::task::spawn_blocking`]
607/// from async contexts. See [`inspect_safetensors_local`] for rationale.
608///
609/// # Errors
610///
611/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
612/// Returns [`FetchError::Io`] if the cached file cannot be read.
613/// Returns [`FetchError::SafetensorsHeader`] if the header is malformed.
614pub fn inspect_safetensors_cached(
615    repo_id: &str,
616    filename: &str,
617    revision: Option<&str>,
618) -> Result<SafetensorsHeaderInfo, FetchError> {
619    let rev = revision.unwrap_or("main");
620
621    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
622        FetchError::SafetensorsHeader {
623            filename: filename.to_owned(),
624            reason: format!("file not found in local cache for {repo_id} ({rev})"),
625        }
626    })?;
627
628    inspect_safetensors_local(&cached_path)
629}
630
631/// Inspects a `.gguf` file's metadata from the local `HuggingFace` cache.
632///
633/// Delegates to [`anamnesis::parse_gguf`] for the on-disk parse, then maps the
634/// result into the format-agnostic [`SafetensorsHeaderInfo`] shape used by
635/// hf-fm's existing render path. Tensor names, GGUF-native shape order, and
636/// dtype name strings carry over directly; per-tensor `data_offsets` are
637/// `(data_offset, data_offset + byte_len)` (with `byte_len = 0` for tensors
638/// whose dtype has no known byte size in anamnesis yet).
639///
640/// **Naming note:** the returned type is still called [`SafetensorsHeaderInfo`]
641/// in v0.10.x because renaming a public type is a breaking change; the
642/// uniform-dispatch rename to a format-agnostic name is scheduled for v0.10.3
643/// when the dispatcher extends across `.npz` / `.pth` (see the cache-management
644/// roadmap). For now, treat the type name as "header / file-level inspect
645/// info" regardless of format.
646///
647/// **Metadata surfacing:** the GGUF metadata table can contain very large
648/// arrays (e.g. tokenizer.ggml.tokens with 50K+ entries). To keep `Metadata:`
649/// rendering useful, this function surfaces *scalar* metadata values only —
650/// strings, booleans, integers, floats — and skips arrays. The GGUF format
651/// version is surfaced under the synthetic key `gguf.version`, the effective
652/// alignment under `gguf.alignment`. The original `general.architecture`,
653/// `general.name`, and friends pass through unchanged.
654///
655/// **Blocking I/O:** anamnesis's GGUF parser mmaps the file; this function is
656/// synchronous and should be wrapped in [`tokio::task::spawn_blocking`] from
657/// async contexts.
658///
659/// # Errors
660///
661/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
662/// Returns [`FetchError::SafetensorsHeader`] if anamnesis rejects the GGUF
663/// file (malformed header, truncated tensor table, etc.).
664pub fn inspect_gguf_cached(
665    repo_id: &str,
666    filename: &str,
667    revision: Option<&str>,
668) -> Result<SafetensorsHeaderInfo, FetchError> {
669    let rev = revision.unwrap_or("main");
670
671    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
672        FetchError::SafetensorsHeader {
673            // BORROW: explicit .to_owned() for owned String in the error variant
674            filename: filename.to_owned(),
675            reason: format!("file not found in local cache for {repo_id} ({rev})"),
676        }
677    })?;
678
679    let file_size = std::fs::metadata(&cached_path).ok().map(|m| m.len());
680
681    let parsed =
682        anamnesis::parse_gguf(&cached_path).map_err(|e| FetchError::SafetensorsHeader {
683            // BORROW: explicit .to_owned() for owned String in the error variant
684            filename: filename.to_owned(),
685            reason: format!("failed to parse GGUF: {e}"),
686        })?;
687
688    let tensors: Vec<TensorInfo> = parsed
689        .tensor_info()
690        .iter()
691        .map(|info| {
692            let start = info.data_offset;
693            let end = info.byte_len.map_or(start, |b| start.saturating_add(b));
694            TensorInfo {
695                // BORROW: explicit .clone() / .to_string() to materialise owned
696                // String + Vec<usize> from anamnesis's borrowed metadata
697                name: info.name.clone(),
698                dtype: info.dtype.to_string(),
699                shape: info.shape.clone(),
700                data_offsets: (start, end),
701            }
702        })
703        .collect();
704
705    // Stringify scalar metadata only; skip arrays (potentially huge — e.g.
706    // tokenizer.ggml.tokens). Add synthetic keys for the format version and
707    // alignment so they appear in the `Metadata:` block.
708    let mut metadata: HashMap<String, String> = parsed
709        .metadata()
710        .iter()
711        // BORROW: explicit .clone() to materialise an owned String key from
712        // the borrowed HashMap iteration
713        .filter_map(|(k, v)| stringify_gguf_scalar(v).map(|s| (k.clone(), s)))
714        .collect();
715    // BORROW: explicit .to_owned() for owned String keys
716    metadata.insert("gguf.version".to_owned(), parsed.version().to_string());
717    metadata.insert("gguf.alignment".to_owned(), parsed.alignment().to_string());
718
719    Ok(SafetensorsHeaderInfo {
720        tensors,
721        metadata: Some(metadata),
722        // GGUF has no discrete "header size" like safetensors's
723        // u64-length-prefix + JSON. The value is left at 0 here; consumers
724        // that care can derive an approximation from `file_size` minus the
725        // tensor byte sum. The `Metadata:` block's `gguf.version` /
726        // `gguf.alignment` keys surface the equivalent format-level info.
727        header_size: 0,
728        file_size,
729        // GGUF quant info (Q4_K_M etc.) is implicit in per-tensor dtypes;
730        // the v0.10.3 Phase C `Format:` / `Size:` lines are safetensors-only.
731        quant_info: None,
732    })
733}
734
735/// Inspects a `.npz` file's metadata from the local `HuggingFace` cache.
736///
737/// Delegates to [`anamnesis::inspect_npz`] for the on-disk parse (which
738/// reads only the ZIP central directory + per-entry NPY headers — no
739/// tensor data), then maps the result into the format-agnostic
740/// [`SafetensorsHeaderInfo`] shape used by hf-fm's existing render path.
741///
742/// **Synthesised offsets.** Anamnesis exposes per-tensor `byte_len` but
743/// not on-disk byte offsets (NPZ tensors live inside a ZIP archive;
744/// offsets are not part of the inspect surface). hf-fm's
745/// `TensorInfo::data_offsets` is synthesised as cumulative `(start, end)`
746/// pairs — `start = sum of previous byte_lens` — so `byte_len()`
747/// (= `end - start`) renders the actual storage size. The synthetic
748/// offsets are NOT on-disk truth; consumers that need real offsets
749/// should defer to v0.11.x once remote-inspect lands.
750///
751/// **Metadata.** `metadata: None` — NPZ has no metadata block analogous
752/// to safetensors's `__metadata__` or GGUF's KV table.
753///
754/// **Header size.** Always `0` — NPZ has no discrete header analogous
755/// to safetensors's `u64`-length-prefix + JSON. Mirrors the GGUF convention.
756///
757/// **Blocking I/O:** anamnesis's NPZ parser opens the file with
758/// `std::fs::File`; this function is synchronous and should be wrapped
759/// in [`tokio::task::spawn_blocking`] from async contexts.
760///
761/// # Errors
762///
763/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
764/// Returns [`FetchError::SafetensorsHeader`] if anamnesis rejects the NPZ
765/// file (malformed ZIP central directory, unsupported NPY dtype, etc.).
766pub fn inspect_npz_cached(
767    repo_id: &str,
768    filename: &str,
769    revision: Option<&str>,
770) -> Result<SafetensorsHeaderInfo, FetchError> {
771    let rev = revision.unwrap_or("main");
772
773    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
774        FetchError::SafetensorsHeader {
775            // BORROW: explicit .to_owned() for owned String in the error variant
776            filename: filename.to_owned(),
777            reason: format!("file not found in local cache for {repo_id} ({rev})"),
778        }
779    })?;
780
781    let file_size = std::fs::metadata(&cached_path).ok().map(|m| m.len());
782
783    let parsed =
784        anamnesis::inspect_npz(&cached_path).map_err(|e| FetchError::SafetensorsHeader {
785            // BORROW: explicit .to_owned() for owned String in the error variant
786            filename: filename.to_owned(),
787            reason: format!("failed to parse NPZ: {e}"),
788        })?;
789
790    let mut tensors: Vec<TensorInfo> = Vec::with_capacity(parsed.tensors.len());
791    let mut cursor: u64 = 0;
792    for t in parsed.tensors {
793        // CAST: usize → u64, byte_len fits in u64 by definition (in-memory size).
794        #[allow(clippy::as_conversions)]
795        let len = t.byte_len as u64;
796        let start = cursor;
797        let end = cursor.saturating_add(len);
798        cursor = end;
799        tensors.push(TensorInfo {
800            name: t.name,
801            // BORROW: explicit .to_string() — anamnesis `NpzDtype` enum → owned `String`
802            dtype: t.dtype.to_string(),
803            shape: t.shape,
804            data_offsets: (start, end),
805        });
806    }
807
808    Ok(SafetensorsHeaderInfo {
809        tensors,
810        metadata: None,
811        header_size: 0,
812        file_size,
813        // NPZ has no quant-method metadata; quant_info stays None.
814        quant_info: None,
815    })
816}
817
818/// Inspects a `.pth` file's metadata from the local `HuggingFace` cache.
819///
820/// Delegates to [`anamnesis::parse_pth`] for the on-disk parse, then uses
821/// the metadata-only `ParsedPth::tensor_info()` view (new in anamnesis
822/// `0.5.0`) to enumerate `(name, shape, dtype, byte_len)` per tensor — no
823/// further I/O beyond the initial mmap. The earlier `.tensors()` method
824/// would materialise each tensor's data via `Cow<'a, [u8]>`, which is
825/// unnecessary for inspect-only use.
826///
827/// **Synthesised offsets.** As with NPZ, anamnesis exposes per-tensor
828/// `byte_len` but not on-disk byte offsets (PTH tensors live inside a
829/// ZIP archive; offsets are not part of the inspect surface). hf-fm's
830/// `TensorInfo::data_offsets` is synthesised as cumulative `(start, end)`
831/// pairs so `byte_len()` (= `end - start`) renders the actual storage size.
832///
833/// **Metadata.** `metadata: None` — PTH has no metadata block analogous
834/// to safetensors's `__metadata__` or GGUF's KV table. The format-level
835/// `big_endian` flag (rare, near-always `false`) is not surfaced here;
836/// can be added as a synthetic `pth.big_endian` key in a future patch if
837/// real users request it.
838///
839/// **Header size.** Always `0` — PTH has no discrete header analogous
840/// to safetensors's `u64`-length-prefix + JSON. Mirrors the GGUF / NPZ
841/// convention.
842///
843/// **Blocking I/O:** anamnesis's PTH parser mmaps the file; this function
844/// is synchronous and should be wrapped in [`tokio::task::spawn_blocking`]
845/// from async contexts.
846///
847/// # Errors
848///
849/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
850/// Returns [`FetchError::SafetensorsHeader`] if anamnesis rejects the PTH
851/// file (malformed pickle stream, legacy pre-1.6 raw-pickle format,
852/// unsupported tensor dtype, etc.).
853pub fn inspect_pth_cached(
854    repo_id: &str,
855    filename: &str,
856    revision: Option<&str>,
857) -> Result<SafetensorsHeaderInfo, FetchError> {
858    let rev = revision.unwrap_or("main");
859
860    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
861        FetchError::SafetensorsHeader {
862            // BORROW: explicit .to_owned() for owned String in the error variant
863            filename: filename.to_owned(),
864            reason: format!("file not found in local cache for {repo_id} ({rev})"),
865        }
866    })?;
867
868    let file_size = std::fs::metadata(&cached_path).ok().map(|m| m.len());
869
870    let parsed = anamnesis::parse_pth(&cached_path).map_err(|e| FetchError::SafetensorsHeader {
871        // BORROW: explicit .to_owned() for owned String in the error variant
872        filename: filename.to_owned(),
873        reason: format!("failed to parse PTH: {e}"),
874    })?;
875
876    let pth_tensors = parsed.tensor_info();
877    let mut tensors: Vec<TensorInfo> = Vec::with_capacity(pth_tensors.len());
878    let mut cursor: u64 = 0;
879    for t in pth_tensors {
880        // CAST: usize → u64, byte_len fits in u64 by definition (in-memory size).
881        #[allow(clippy::as_conversions)]
882        let len = t.byte_len as u64;
883        let start = cursor;
884        let end = cursor.saturating_add(len);
885        cursor = end;
886        tensors.push(TensorInfo {
887            name: t.name,
888            // BORROW: explicit .to_string() — anamnesis `PthDtype` enum → owned `String`
889            dtype: t.dtype.to_string(),
890            shape: t.shape,
891            data_offsets: (start, end),
892        });
893    }
894
895    Ok(SafetensorsHeaderInfo {
896        tensors,
897        metadata: None,
898        header_size: 0,
899        file_size,
900        // PTH has no quant-method metadata; quant_info stays None.
901        quant_info: None,
902    })
903}
904
905/// Stringifies a scalar `GgufMetadataValue` from anamnesis.
906///
907/// Returns `None` for array variants (potentially huge — vocab tables, merges
908/// lists) and for any future `#[non_exhaustive]` variants we don't yet
909/// recognise. Surfaced through the `Metadata:` block in `inspect` output by
910/// [`inspect_gguf_cached`].
911//
912// `GgufMetadataValue` is `#[non_exhaustive]`. The explicit `V::Array(_)` arm
913// and the `_ =>` catch-all both return `None`, but they document different
914// intents — "array variants are deliberately skipped" vs "future unknown
915// variants fall through". Clippy's `match_same_arms` flags the bodies as
916// identical; the duplication is intentional.
917#[allow(clippy::match_same_arms)]
918fn stringify_gguf_scalar(value: &anamnesis::parse::gguf::GgufMetadataValue) -> Option<String> {
919    use anamnesis::parse::gguf::GgufMetadataValue as V;
920    match value {
921        V::String(s) => Some(s.clone()),
922        V::Bool(b) => Some(b.to_string()),
923        V::U8(n) => Some(n.to_string()),
924        V::I8(n) => Some(n.to_string()),
925        V::U16(n) => Some(n.to_string()),
926        V::I16(n) => Some(n.to_string()),
927        V::U32(n) => Some(n.to_string()),
928        V::I32(n) => Some(n.to_string()),
929        V::U64(n) => Some(n.to_string()),
930        V::I64(n) => Some(n.to_string()),
931        V::F32(n) => Some(format!("{n}")),
932        V::F64(n) => Some(format!("{n}")),
933        V::Array(_) => None,
934        _ => None,
935    }
936}
937
938// -----------------------------------------------------------------------
939// Public API: multi-file inspection
940// -----------------------------------------------------------------------
941
942/// Inspects all `.safetensors` files in a repository (cache-first per file).
943///
944/// Fetches the file listing via `list_repo_files_with_metadata()`, then
945/// inspects each `.safetensors` file's header via [`inspect_safetensors()`].
946/// For each file, checks the local cache first and only makes HTTP Range
947/// requests on cache miss. Returns full per-shard headers in filename order.
948///
949/// For a lightweight summary of sharded models (tensor counts per shard
950/// without fetching individual headers), use [`fetch_shard_index()`] instead.
951///
952/// # Errors
953///
954/// Returns [`FetchError::Http`] if the metadata or Range requests fail.
955pub async fn inspect_repo_safetensors(
956    repo_id: &str,
957    token: Option<&str>,
958    revision: Option<&str>,
959) -> Result<Vec<(String, SafetensorsHeaderInfo, InspectSource)>, FetchError> {
960    let client = crate::chunked::build_client(token)?;
961    let files =
962        crate::repo::list_repo_files_with_metadata(repo_id, token, revision, &client).await?;
963
964    let safetensors_files: Vec<String> = files
965        .into_iter()
966        .filter(|f| f.filename.ends_with(".safetensors"))
967        .map(|f| f.filename)
968        .collect();
969
970    if safetensors_files.is_empty() {
971        return Ok(Vec::new());
972    }
973
974    let semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new(4));
975    let mut join_set = JoinSet::new();
976
977    for filename in safetensors_files {
978        // BORROW: explicit .clone()/.to_owned() to move into async task
979        let sem = semaphore.clone();
980        let repo = repo_id.to_owned();
981        let tok = token.map(str::to_owned);
982        let rev = revision.map(str::to_owned);
983
984        join_set.spawn(async move {
985            let _permit = sem
986                .acquire()
987                .await
988                .map_err(|e| FetchError::Http(format!("semaphore error: {e}")))?;
989            // BORROW: explicit .as_deref() for Option<String> → Option<&str>
990            let (info, source) =
991                inspect_safetensors(&repo, &filename, tok.as_deref(), rev.as_deref()).await?;
992            Ok::<_, FetchError>((filename, info, source))
993        });
994    }
995
996    let mut results = Vec::new();
997    while let Some(join_result) = join_set.join_next().await {
998        match join_result {
999            Ok(Ok(item)) => results.push(item),
1000            Ok(Err(e)) => {
1001                join_set.abort_all();
1002                return Err(e);
1003            }
1004            Err(e) => {
1005                join_set.abort_all();
1006                return Err(FetchError::Http(format!("task join error: {e}")));
1007            }
1008        }
1009    }
1010
1011    results.sort_by(|a, b| a.0.cmp(&b.0));
1012
1013    Ok(results)
1014}
1015
1016/// A `(filename, size_bytes)` enumeration of safetensors files in a repo,
1017/// paired with the commit SHA of the resolved revision (when known).
1018///
1019/// The same tuple shape serves both local and remote listings: [`list_cached_safetensors`]
1020/// produces it from a cached snapshot; `repo::list_repo_files_with_commit` filtered to
1021/// `*.safetensors` produces it from the `HuggingFace` API. Callers that need a uniform
1022/// view over "what safetensors can I inspect?" regardless of source use this alias.
1023pub type SafetensorsListing = (Vec<(String, u64)>, Option<String>);
1024
1025/// Lists `.safetensors` files in the cached snapshot for `repo_id`@`revision`.
1026///
1027/// Returns `(entries, commit_sha)` where `entries` is a sorted list of
1028/// `(filename, size_bytes)` tuples, and `commit_sha` is the snapshot's commit
1029/// hash (same value stored in `refs/<revision>`). Returns empty lists when the
1030/// repo or revision is not cached. Unlike [`inspect_repo_safetensors_cached`],
1031/// this does **not** parse any headers — it is a cheap name-and-size enumeration
1032/// intended for discovery UI (e.g. `inspect --list --cached`).
1033///
1034/// # Blocking I/O
1035///
1036/// Performs a synchronous recursive directory walk with a `stat` call per
1037/// `.safetensors` entry. On local SSDs the cost is sub-millisecond; on
1038/// networked caches (NFS/CIFS) a large sharded repo can take seconds. Wrap
1039/// in [`tokio::task::spawn_blocking`] from async contexts.
1040///
1041/// # Errors
1042///
1043/// Returns [`FetchError::Io`] if the snapshot directory cannot be read.
1044pub fn list_cached_safetensors(
1045    repo_id: &str,
1046    revision: Option<&str>,
1047) -> Result<SafetensorsListing, FetchError> {
1048    let rev = revision.unwrap_or("main");
1049    let cache_dir = cache::hf_cache_dir()?;
1050    let repo_dir = cache_layout::repo_dir(&cache_dir, repo_id);
1051
1052    let Some(commit_hash) = cache::read_ref(&repo_dir, rev) else {
1053        return Ok((Vec::new(), None));
1054    };
1055
1056    let snapshot_dir = cache_layout::snapshot_dir(&repo_dir, &commit_hash);
1057    if !snapshot_dir.exists() {
1058        return Ok((Vec::new(), Some(commit_hash)));
1059    }
1060
1061    let mut results = Vec::new();
1062    collect_safetensors_names_sizes(&snapshot_dir, "", &mut results)?;
1063    results.sort_by(|a, b| a.0.cmp(&b.0));
1064    Ok((results, Some(commit_hash)))
1065}
1066
1067/// Recursively collects `(filename, size)` pairs for `.safetensors` files.
1068fn collect_safetensors_names_sizes(
1069    dir: &Path,
1070    prefix: &str,
1071    results: &mut Vec<(String, u64)>,
1072) -> Result<(), FetchError> {
1073    let entries = std::fs::read_dir(dir).map_err(|e| FetchError::Io {
1074        path: dir.to_path_buf(),
1075        source: e,
1076    })?;
1077
1078    for entry in entries {
1079        let Ok(entry) = entry else { continue };
1080        let path = entry.path();
1081        // BORROW: explicit .to_string_lossy() for OsString → str conversion
1082        let name = entry.file_name().to_string_lossy().to_string();
1083
1084        if path.is_dir() {
1085            let child_prefix = if prefix.is_empty() {
1086                name
1087            } else {
1088                format!("{prefix}/{name}")
1089            };
1090            collect_safetensors_names_sizes(&path, &child_prefix, results)?;
1091        } else if name.ends_with(".safetensors") {
1092            let filename = if prefix.is_empty() {
1093                name
1094            } else {
1095                format!("{prefix}/{name}")
1096            };
1097            let size = entry.metadata().map_or(0, |m| m.len());
1098            results.push((filename, size));
1099        }
1100    }
1101
1102    Ok(())
1103}
1104
1105/// Inspects all `.safetensors` files in a cached repository (no network).
1106///
1107/// Walks the snapshot directory and inspects each `.safetensors` file's
1108/// header from local disk. Returns results in filename order.
1109///
1110/// # Blocking I/O
1111///
1112/// Walks the snapshot directory and reads each header synchronously. In async
1113/// contexts, wrap in [`tokio::task::spawn_blocking`] to avoid stalling the
1114/// runtime — multi-shard repos on network-mounted caches can take seconds.
1115///
1116/// # Errors
1117///
1118/// Returns [`FetchError::Io`] if the cache directory cannot be read.
1119/// Returns [`FetchError::SafetensorsHeader`] if any header is malformed.
1120pub fn inspect_repo_safetensors_cached(
1121    repo_id: &str,
1122    revision: Option<&str>,
1123) -> Result<Vec<(String, SafetensorsHeaderInfo)>, FetchError> {
1124    let rev = revision.unwrap_or("main");
1125    let cache_dir = cache::hf_cache_dir()?;
1126    let repo_dir = cache_layout::repo_dir(&cache_dir, repo_id);
1127
1128    let Some(commit_hash) = cache::read_ref(&repo_dir, rev) else {
1129        return Ok(Vec::new());
1130    };
1131
1132    let snapshot_dir = cache_layout::snapshot_dir(&repo_dir, &commit_hash);
1133    if !snapshot_dir.exists() {
1134        return Ok(Vec::new());
1135    }
1136
1137    let mut results = Vec::new();
1138    collect_safetensors_recursive(&snapshot_dir, "", &mut results)?;
1139    results.sort_by(|a, b| a.0.cmp(&b.0));
1140
1141    Ok(results)
1142}
1143
1144/// Recursively finds and inspects `.safetensors` files in a snapshot directory.
1145fn collect_safetensors_recursive(
1146    dir: &Path,
1147    prefix: &str,
1148    results: &mut Vec<(String, SafetensorsHeaderInfo)>,
1149) -> Result<(), FetchError> {
1150    let entries = std::fs::read_dir(dir).map_err(|e| FetchError::Io {
1151        path: dir.to_path_buf(),
1152        source: e,
1153    })?;
1154
1155    for entry in entries {
1156        let Ok(entry) = entry else { continue };
1157        let path = entry.path();
1158        // BORROW: explicit .to_string_lossy() for OsString → str conversion
1159        let name = entry.file_name().to_string_lossy().to_string();
1160
1161        if path.is_dir() {
1162            let child_prefix = if prefix.is_empty() {
1163                name
1164            } else {
1165                format!("{prefix}/{name}")
1166            };
1167            collect_safetensors_recursive(&path, &child_prefix, results)?;
1168        } else if name.ends_with(".safetensors") {
1169            let filename = if prefix.is_empty() {
1170                name
1171            } else {
1172                format!("{prefix}/{name}")
1173            };
1174            let info = inspect_safetensors_local(&path)?;
1175            results.push((filename, info));
1176        }
1177    }
1178
1179    Ok(())
1180}
1181
1182// -----------------------------------------------------------------------
1183// Shard index
1184// -----------------------------------------------------------------------
1185
1186/// Raw JSON structure of `model.safetensors.index.json`.
1187#[derive(serde::Deserialize)]
1188struct RawShardIndex {
1189    weight_map: HashMap<String, String>,
1190    #[serde(default)]
1191    metadata: Option<HashMap<String, serde_json::Value>>,
1192}
1193
1194/// Fetches and parses the shard index for a sharded `.safetensors` model (cache-first).
1195///
1196/// Returns `Ok(None)` if the repo has no `model.safetensors.index.json` (i.e.,
1197/// the model is not sharded or uses a single `.safetensors` file).
1198///
1199/// # Errors
1200///
1201/// Returns [`FetchError::Http`] if the index fetch fails.
1202/// Returns [`FetchError::SafetensorsHeader`] if the index JSON is malformed.
1203pub async fn fetch_shard_index(
1204    repo_id: &str,
1205    token: Option<&str>,
1206    revision: Option<&str>,
1207) -> Result<Option<ShardedIndex>, FetchError> {
1208    let rev = revision.unwrap_or("main");
1209    let index_filename = "model.safetensors.index.json";
1210
1211    // Try local cache first.
1212    if let Some(cached_path) = resolve_cached_path(repo_id, rev, index_filename) {
1213        let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
1214            path: cached_path,
1215            source: e,
1216        })?;
1217        let index = parse_shard_index_json(&content, repo_id)?;
1218        return Ok(Some(index));
1219    }
1220
1221    // Fall back to HTTP.
1222    let client = chunked::build_client(token)?;
1223    let url = chunked::build_download_url(repo_id, rev, index_filename);
1224
1225    // BORROW: explicit .as_str() instead of Deref coercion
1226    let response =
1227        client.get(url.as_str()).send().await.map_err(|e| {
1228            FetchError::Http(format!("failed to fetch shard index for {repo_id}: {e}"))
1229        })?;
1230
1231    if response.status() == reqwest::StatusCode::NOT_FOUND {
1232        return Ok(None);
1233    }
1234
1235    if !response.status().is_success() {
1236        return Err(FetchError::Http(format!(
1237            "shard index request for {repo_id} returned status {}",
1238            response.status()
1239        )));
1240    }
1241
1242    let content = response
1243        .text()
1244        .await
1245        .map_err(|e| FetchError::Http(format!("failed to read shard index for {repo_id}: {e}")))?;
1246
1247    let index = parse_shard_index_json(&content, repo_id)?;
1248    Ok(Some(index))
1249}
1250
1251/// Fetches the shard index from cache only (no network).
1252///
1253/// Returns `Ok(None)` if the index file is not cached.
1254///
1255/// # Errors
1256///
1257/// Returns [`FetchError::Io`] if the cached file cannot be read.
1258/// Returns [`FetchError::SafetensorsHeader`] if the index JSON is malformed.
1259pub fn fetch_shard_index_cached(
1260    repo_id: &str,
1261    revision: Option<&str>,
1262) -> Result<Option<ShardedIndex>, FetchError> {
1263    let rev = revision.unwrap_or("main");
1264    let index_filename = "model.safetensors.index.json";
1265
1266    let Some(cached_path) = resolve_cached_path(repo_id, rev, index_filename) else {
1267        return Ok(None);
1268    };
1269
1270    let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
1271        path: cached_path,
1272        source: e,
1273    })?;
1274
1275    let index = parse_shard_index_json(&content, repo_id)?;
1276    Ok(Some(index))
1277}
1278
1279/// Parses shard index JSON into a `ShardedIndex`.
1280fn parse_shard_index_json(content: &str, repo_id: &str) -> Result<ShardedIndex, FetchError> {
1281    let raw: RawShardIndex =
1282        serde_json::from_str(content).map_err(|e| FetchError::SafetensorsHeader {
1283            filename: "model.safetensors.index.json".to_owned(),
1284            reason: format!("failed to parse shard index for {repo_id}: {e}"),
1285        })?;
1286
1287    // Collect unique shard filenames in sorted order.
1288    let mut shard_set: Vec<String> = raw.weight_map.values().cloned().collect();
1289    shard_set.sort();
1290    shard_set.dedup();
1291
1292    Ok(ShardedIndex {
1293        weight_map: raw.weight_map,
1294        shards: shard_set,
1295        metadata: raw.metadata,
1296    })
1297}
1298
1299// -----------------------------------------------------------------------
1300// Param formatting helper
1301// -----------------------------------------------------------------------
1302
1303/// Formats a parameter count with a compact suffix (e.g., `927.0M`, `1.02B`).
1304#[must_use]
1305pub fn format_params(count: u64) -> String {
1306    // CAST: u64 → f64, precision loss acceptable; value is a display-only scalar
1307    #[allow(clippy::cast_precision_loss, clippy::as_conversions)]
1308    let val = count as f64;
1309
1310    if count >= 1_000_000_000 {
1311        format!("{:.2}B", val / 1_000_000_000.0)
1312    } else if count >= 1_000_000 {
1313        format!("{:.1}M", val / 1_000_000.0)
1314    } else if count >= 1_000 {
1315        format!("{:.1}K", val / 1_000.0)
1316    } else {
1317        count.to_string()
1318    }
1319}
1320
1321// -----------------------------------------------------------------------
1322// Adapter config
1323// -----------------------------------------------------------------------
1324
1325/// Raw JSON structure of `adapter_config.json`.
1326#[derive(serde::Deserialize)]
1327struct RawAdapterConfig {
1328    #[serde(default)]
1329    peft_type: Option<String>,
1330    #[serde(default)]
1331    base_model_name_or_path: Option<String>,
1332    #[serde(default)]
1333    r: Option<u32>,
1334    #[serde(default)]
1335    lora_alpha: Option<f64>,
1336    #[serde(default)]
1337    target_modules: Option<AdapterTargetModules>,
1338    #[serde(default)]
1339    task_type: Option<String>,
1340}
1341
1342/// `target_modules` in adapter configs can be a list of strings or a single string.
1343#[derive(serde::Deserialize)]
1344#[serde(untagged)]
1345enum AdapterTargetModules {
1346    /// A list of module name strings.
1347    List(Vec<String>),
1348    /// A single module name string.
1349    Single(String),
1350}
1351
1352/// Fetches and parses `adapter_config.json` for a `PEFT` adapter repository (cache-first).
1353///
1354/// Returns `Ok(None)` if the file does not exist (HTTP 404), meaning the
1355/// repository is not a `PEFT` adapter.
1356///
1357/// # Errors
1358///
1359/// Returns [`FetchError::Http`] if the request fails (other than 404).
1360/// Returns [`FetchError::SafetensorsHeader`] if the JSON is malformed.
1361pub async fn fetch_adapter_config(
1362    repo_id: &str,
1363    token: Option<&str>,
1364    revision: Option<&str>,
1365) -> Result<Option<AdapterConfig>, FetchError> {
1366    let rev = revision.unwrap_or("main");
1367    let config_filename = "adapter_config.json";
1368
1369    // Try local cache first.
1370    if let Some(cached_path) = resolve_cached_path(repo_id, rev, config_filename) {
1371        let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
1372            path: cached_path,
1373            source: e,
1374        })?;
1375        let config = parse_adapter_config_json(&content, repo_id)?;
1376        return Ok(Some(config));
1377    }
1378
1379    // Fall back to HTTP.
1380    let client = chunked::build_client(token)?;
1381    let url = chunked::build_download_url(repo_id, rev, config_filename);
1382
1383    // BORROW: explicit .as_str() instead of Deref coercion
1384    let response = client.get(url.as_str()).send().await.map_err(|e| {
1385        FetchError::Http(format!("failed to fetch adapter config for {repo_id}: {e}"))
1386    })?;
1387
1388    if response.status() == reqwest::StatusCode::NOT_FOUND {
1389        return Ok(None);
1390    }
1391
1392    if !response.status().is_success() {
1393        return Err(FetchError::Http(format!(
1394            "adapter config request for {repo_id} returned status {}",
1395            response.status()
1396        )));
1397    }
1398
1399    let content = response.text().await.map_err(|e| {
1400        FetchError::Http(format!("failed to read adapter config for {repo_id}: {e}"))
1401    })?;
1402
1403    let config = parse_adapter_config_json(&content, repo_id)?;
1404    Ok(Some(config))
1405}
1406
1407/// Fetches the adapter config from cache only (no network).
1408///
1409/// Returns `Ok(None)` if the file is not cached.
1410///
1411/// # Errors
1412///
1413/// Returns [`FetchError::Io`] if the cached file cannot be read.
1414/// Returns [`FetchError::SafetensorsHeader`] if the JSON is malformed.
1415pub fn fetch_adapter_config_cached(
1416    repo_id: &str,
1417    revision: Option<&str>,
1418) -> Result<Option<AdapterConfig>, FetchError> {
1419    let rev = revision.unwrap_or("main");
1420    let config_filename = "adapter_config.json";
1421
1422    let Some(cached_path) = resolve_cached_path(repo_id, rev, config_filename) else {
1423        return Ok(None);
1424    };
1425
1426    let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
1427        path: cached_path,
1428        source: e,
1429    })?;
1430
1431    let config = parse_adapter_config_json(&content, repo_id)?;
1432    Ok(Some(config))
1433}
1434
1435/// Parses adapter config JSON into an [`AdapterConfig`].
1436fn parse_adapter_config_json(content: &str, repo_id: &str) -> Result<AdapterConfig, FetchError> {
1437    let raw: RawAdapterConfig =
1438        serde_json::from_str(content).map_err(|e| FetchError::SafetensorsHeader {
1439            filename: "adapter_config.json".to_owned(),
1440            reason: format!("failed to parse adapter config for {repo_id}: {e}"),
1441        })?;
1442
1443    let target_modules = match raw.target_modules {
1444        Some(AdapterTargetModules::List(v)) => v,
1445        Some(AdapterTargetModules::Single(s)) => vec![s],
1446        None => Vec::new(),
1447    };
1448
1449    Ok(AdapterConfig {
1450        peft_type: raw.peft_type,
1451        base_model_name_or_path: raw.base_model_name_or_path,
1452        r: raw.r,
1453        lora_alpha: raw.lora_alpha,
1454        target_modules,
1455        task_type: raw.task_type,
1456    })
1457}