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}