ferritin_plms/
loader.rs

1//! One place to describe where a model's weights live and how to load them.
2//!
3//! Before this module every runner reimplemented the same download block:
4//! split the repo id, build an [`HFClientSync`], download a file, then either
5//! `unsafe { VarBuilder::from_mmaped_safetensors(..) }` or
6//! `PthTensors::new(..)`. Six near-identical copies drifted apart — only ESM3
7//! attached any error context, and every one of them inherited the same
8//! silently-wrong repo-id split (`split_once('/').unwrap_or(("", repo_id))`,
9//! which turns a malformed id into an empty owner and a confusing 404 rather
10//! than a clear error).
11//!
12//! The pieces here are:
13//!
14//! - [`WeightSource`] — plain `const` data on each model enum: which repo,
15//!   which revision, and which on-disk [`Format`] the weights use.
16//! - [`LoadOptions`] — the device and dtype to load onto, so the six
17//!   per-module `const *_DTYPE` definitions collapse into one field.
18//! - [`WeightSource::var_builder`] — the single place holding the `unsafe`
19//!   mmap.
20//! - [`optional_prefix`] — the "is this checkpoint wrapped in an HF
21//!   `*ForMaskedLM` class?" probe, generalised from ESMC's `esmc.` special
22//!   case.
23//!
24//! ```no_run
25//! # use ferritin_plms::loader::{LoadOptions, WeightSource};
26//! # use candle_core::Device;
27//! const WEIGHTS: WeightSource = WeightSource::safetensors("facebook/esm2_t6_8M_UR50D");
28//!
29//! let opts = LoadOptions::new(Device::Cpu);
30//! let vb = WEIGHTS.var_builder("model.safetensors", &opts)?;
31//! # Ok::<(), anyhow::Error>(())
32//! ```
33
34use anyhow::{Context, Result, bail};
35use candle_core::pickle::PthTensors;
36use candle_core::{DType, Device};
37use candle_nn::VarBuilder;
38use hf_hub::{HFClientSync, HFError, HFRepositorySync, RepoTypeModel};
39use std::path::{Path, PathBuf};
40use std::time::{Duration, Instant};
41
42/// How long to keep retrying a download that is blocked on another
43/// downloader's cache lock.
44///
45/// hf-hub's own lock timeout is a hardcoded 10 seconds
46/// (`CACHE_LOCK_TIMEOUT_SECS`), which is far shorter than a multi-hundred-MB
47/// download takes on a cold cache. So whenever two threads or processes want
48/// the same not-yet-cached file, the one that does not win the lock fails
49/// almost immediately (ferritin-100.25).
50const CACHE_LOCK_RETRY_BUDGET: Duration = Duration::from_secs(600);
51
52/// Pause between retries while waiting for another downloader to finish.
53const CACHE_LOCK_RETRY_PAUSE: Duration = Duration::from_secs(2);
54
55// ── LoadOptions ───────────────────────────────────────────────────────────────
56
57/// Device and dtype to load a model onto.
58///
59/// Defaults to `F32`, which is what every runner used before this existed.
60#[derive(Debug, Clone)]
61pub struct LoadOptions {
62    /// Element type to materialise weights as.
63    pub dtype: DType,
64    /// Device to place weights on.
65    pub device: Device,
66}
67
68impl LoadOptions {
69    /// `F32` on `device`.
70    pub fn new(device: Device) -> Self {
71        Self {
72            dtype: DType::F32,
73            device,
74        }
75    }
76
77    /// Override the dtype (builder style).
78    pub fn with_dtype(mut self, dtype: DType) -> Self {
79        self.dtype = dtype;
80        self
81    }
82
83    /// Reject dtype/device combinations the backend cannot run.
84    ///
85    /// candle's CPU backend has no BF16 `matmul`, so a BF16 model on CPU fails
86    /// partway through loading or on the first forward pass with a bare
87    /// "unsupported dtype BF16 for op matmul". Catching it here says what to do
88    /// instead (ferritin-100.9).
89    ///
90    /// The refusal is scoped to the CPU backend rather than to the dtype
91    /// because Metal does run BF16: ferritin-100.19 loaded and ran both ESM2
92    /// and AMPLIFY there at BF16 without changing a single top-1 prediction.
93    /// See `tests/test_plm_dtype_parity.rs` for the divergence and throughput
94    /// tables that establish it. CUDA is assumed to behave like Metal here
95    /// and remains unmeasured — the tests cover whatever accelerator the build
96    /// actually has, so running them on a CUDA box is what closes that gap.
97    pub fn validate(&self) -> Result<()> {
98        if self.dtype == DType::BF16 && self.device.is_cpu() {
99            bail!(
100                "BF16 is not supported on the CPU backend: candle has no BF16 matmul there, \
101                 so loading or the first forward pass would fail with a bare \
102                 'unsupported dtype BF16 for op matmul'. Use F16 for half precision on CPU, \
103                 or run on Metal/CUDA."
104            );
105        }
106        Ok(())
107    }
108}
109
110// ── Format ────────────────────────────────────────────────────────────────────
111
112/// How a checkpoint is stored on disk.
113#[derive(Debug, Clone, Copy, PartialEq, Eq)]
114pub enum Format {
115    /// `.safetensors`, loaded by mmap.
116    Safetensors,
117    /// PyTorch `.pth`/`.pt` pickle.
118    Pth {
119        /// Sub-dictionary holding the tensors, when the checkpoint nests them
120        /// (e.g. ProteinMPNN's `model_state_dict`). `None` = tensors at the root.
121        root_key: Option<&'static str>,
122    },
123}
124
125// ── WeightSource ──────────────────────────────────────────────────────────────
126
127/// Where a model's weights live on the HuggingFace hub, and how to read them.
128///
129/// Intended to be a `const` on each model enum:
130///
131/// ```
132/// # use ferritin_plms::loader::WeightSource;
133/// const AMP120M: WeightSource = WeightSource::safetensors("chandar-lab/AMPLIFY_120M");
134/// ```
135#[derive(Debug, Clone, Copy)]
136pub struct WeightSource {
137    /// Full `owner/name` repo id.
138    pub repo_id: &'static str,
139    /// Git revision; `None` means the hub default (`main`).
140    pub revision: Option<&'static str>,
141    /// On-disk format of the weight files.
142    pub format: Format,
143}
144
145impl WeightSource {
146    /// A safetensors repo at the default revision.
147    pub const fn safetensors(repo_id: &'static str) -> Self {
148        Self {
149            repo_id,
150            revision: None,
151            format: Format::Safetensors,
152        }
153    }
154
155    /// A PyTorch-pickle repo at the default revision.
156    ///
157    /// `root_key` names the sub-dictionary holding the tensors, or `None` if
158    /// they sit at the root.
159    pub const fn pth(repo_id: &'static str, root_key: Option<&'static str>) -> Self {
160        Self {
161            repo_id,
162            revision: None,
163            format: Format::Pth { root_key },
164        }
165    }
166
167    /// Pin to a specific revision.
168    pub const fn at_revision(mut self, revision: &'static str) -> Self {
169        self.revision = Some(revision);
170        self
171    }
172
173    /// Split and validate `repo_id` into `(owner, name)`.
174    ///
175    /// Unlike `split_once('/').unwrap_or(("", repo_id))` — which every runner
176    /// used to do — a malformed id is an error here rather than an empty owner
177    /// and a 404 several seconds later.
178    fn owner_and_name(&self) -> Result<(&'static str, &'static str)> {
179        let (owner, name) = self.repo_id.split_once('/').ok_or_else(|| {
180            anyhow::anyhow!(
181                "malformed HuggingFace repo id {:?}: expected 'owner/name'",
182                self.repo_id
183            )
184        })?;
185        if owner.is_empty() || name.is_empty() || name.contains('/') {
186            bail!(
187                "malformed HuggingFace repo id {:?}: expected 'owner/name' with both parts non-empty",
188                self.repo_id
189            );
190        }
191        Ok((owner, name))
192    }
193
194    fn repo(&self) -> Result<HFRepositorySync<RepoTypeModel>> {
195        let (owner, name) = self.owner_and_name()?;
196        let client = HFClientSync::new().with_context(|| {
197            format!(
198                "failed to initialise the HuggingFace client while loading {}",
199                self.repo_id
200            )
201        })?;
202        Ok(client.model(owner, name))
203    }
204
205    /// Download `filename`, returning its cached path.
206    ///
207    /// Errors name both the repo and the file — previously only ESM3 did this,
208    /// so a failure elsewhere surfaced as a bare transport error.
209    pub fn fetch(&self, filename: &str) -> Result<PathBuf> {
210        let repo = self.repo()?;
211        let deadline = Instant::now() + CACHE_LOCK_RETRY_BUDGET;
212
213        loop {
214            let attempt = repo
215                .download_file()
216                .filename(filename.to_string())
217                .maybe_revision(self.revision.map(str::to_string))
218                .send();
219
220            match attempt {
221                Ok(path) => return Ok(path),
222
223                // Another thread or process is downloading this same file and
224                // holds its cache lock. hf-hub gives up after 10 seconds, which
225                // a large download will always exceed — so wait for the
226                // downloader rather than failing (ferritin-100.25).
227                Err(HFError::CacheLockTimeout { .. }) if Instant::now() < deadline => {
228                    eprintln!(
229                        "{}: waiting on another download of {filename} (cache lock held)…",
230                        self.repo_id
231                    );
232                    std::thread::sleep(CACHE_LOCK_RETRY_PAUSE);
233                }
234
235                Err(e) => {
236                    return Err(anyhow::Error::new(e)).with_context(|| {
237                        format!("failed to download {filename} from {}", self.repo_id)
238                    });
239                }
240            }
241        }
242    }
243
244    /// Download `filename`, returning `None` if it is absent or unreachable.
245    ///
246    /// For genuinely optional files such as ESM-2's `config.json`, where the
247    /// runner falls back to a built-in config.
248    pub fn fetch_optional(&self, filename: &str) -> Option<PathBuf> {
249        self.fetch(filename).ok()
250    }
251
252    /// Download `filename` and build a [`VarBuilder`] over it.
253    ///
254    /// This is the only place in the crate that performs the safetensors mmap,
255    /// and the only place the dtype is applied.
256    /// Suffix marking a HuggingFace shard index rather than a weight file.
257    const SHARD_INDEX_SUFFIX: &'static str = ".index.json";
258
259    pub fn var_builder(&self, filename: &str, opts: &LoadOptions) -> Result<VarBuilder<'static>> {
260        if filename.ends_with(Self::SHARD_INDEX_SUFFIX) {
261            return self.var_builder_sharded(filename, opts);
262        }
263        let path = self.fetch(filename)?;
264        self.var_builder_from_path(&path, opts)
265    }
266
267    /// Build a [`VarBuilder`] over a checkpoint split across several
268    /// safetensors shards.
269    ///
270    /// Large checkpoints ship as `model-0000N-of-0000M.safetensors` plus a
271    /// `model.safetensors.index.json` whose `weight_map` says which shard holds
272    /// each tensor. candle's `from_mmaped_safetensors` already takes a slice of
273    /// paths, so all this does is read the index, fetch every shard it names,
274    /// and hand candle the lot (ferritin-100.24).
275    ///
276    /// Shard order does not matter — candle indexes by tensor name — but the
277    /// shards are deduplicated and sorted so the mmap set is deterministic.
278    pub fn var_builder_sharded(
279        &self,
280        index_filename: &str,
281        opts: &LoadOptions,
282    ) -> Result<VarBuilder<'static>> {
283        opts.validate()?;
284        if !matches!(self.format, Format::Safetensors) {
285            bail!(
286                "{}: sharded loading is only defined for safetensors, but {index_filename} \
287                 was requested for a {:?} source",
288                self.repo_id,
289                self.format
290            );
291        }
292
293        let index_path = self.fetch(index_filename)?;
294        let shards = Self::shard_names(&index_path)?;
295
296        let paths = shards
297            .iter()
298            .map(|shard| self.fetch(shard))
299            .collect::<Result<Vec<_>>>()?;
300
301        // SAFETY: as in var_builder_from_path — mmap of files we just
302        // materialised in the HF cache.
303        let vb = unsafe {
304            VarBuilder::from_mmaped_safetensors(&paths, opts.dtype, &opts.device).with_context(
305                || {
306                    format!(
307                        "failed to mmap {} shards from {}",
308                        paths.len(),
309                        self.repo_id
310                    )
311                },
312            )?
313        };
314        Ok(vb)
315    }
316
317    /// Read a shard index and return the shard filenames it references, sorted
318    /// and deduplicated.
319    fn shard_names(index_path: &Path) -> Result<Vec<String>> {
320        #[derive(serde::Deserialize)]
321        struct ShardIndex {
322            weight_map: std::collections::BTreeMap<String, String>,
323        }
324
325        let raw = std::fs::read_to_string(index_path)
326            .with_context(|| format!("failed to read shard index {}", index_path.display()))?;
327        let index: ShardIndex = serde_json::from_str(&raw).with_context(|| {
328            format!(
329                "shard index {} is not a HuggingFace weight_map",
330                index_path.display()
331            )
332        })?;
333
334        if index.weight_map.is_empty() {
335            bail!(
336                "shard index {} has an empty weight_map",
337                index_path.display()
338            );
339        }
340
341        let mut shards: Vec<String> = index.weight_map.into_values().collect();
342        shards.sort_unstable();
343        shards.dedup();
344        Ok(shards)
345    }
346
347    /// Build a [`VarBuilder`] over an already-downloaded (or local) file.
348    ///
349    /// Used by loaders that also accept a path directly, such as
350    /// `ProteinMPNNRunner::from_path`.
351    pub fn var_builder_from_path(
352        &self,
353        path: &Path,
354        opts: &LoadOptions,
355    ) -> Result<VarBuilder<'static>> {
356        var_builder_from_path(path, self.format, opts)
357    }
358}
359
360/// Build a [`VarBuilder`] over a local weight file of the given [`Format`].
361///
362/// Free-standing counterpart to [`WeightSource::var_builder_from_path`], for
363/// paths with no associated hub repo.
364pub fn var_builder_from_path(
365    path: &Path,
366    format: Format,
367    opts: &LoadOptions,
368) -> Result<VarBuilder<'static>> {
369    opts.validate()?;
370    match format {
371        Format::Safetensors => {
372            // SAFETY: mmap of a file we just materialised in the HF cache. As
373            // everywhere in candle, this assumes nothing else mutates the file
374            // while it is mapped.
375            let vb = unsafe {
376                VarBuilder::from_mmaped_safetensors(&[path], opts.dtype, &opts.device)
377                    .with_context(|| format!("failed to mmap {}", path.display()))?
378            };
379            Ok(vb)
380        }
381        Format::Pth { root_key } => {
382            let pth = PthTensors::new(path, root_key)
383                .with_context(|| format!("failed to parse {}", path.display()))?;
384            Ok(VarBuilder::from_backend(
385                Box::new(pth),
386                opts.dtype,
387                opts.device.clone(),
388            ))
389        }
390    }
391}
392
393// ── Optional wrapper prefix ───────────────────────────────────────────────────
394
395/// Descend into `prefix` when the checkpoint nests the backbone under it.
396///
397/// HuggingFace wrapper classes (`EsmForMaskedLM`, `ESMCForMaskedLM`, …) store
398/// the backbone under an attribute, so the same architecture ships both flat
399/// and prefixed. `probe` is a tensor that always exists in the backbone; if
400/// `{prefix}.{probe}` is present the prefixed root is returned, otherwise `vb`
401/// is returned unchanged.
402///
403/// ```no_run
404/// # use ferritin_plms::loader::optional_prefix;
405/// # fn f(vb: candle_nn::VarBuilder<'static>) {
406/// // "esmc.embed.weight" present → root at "esmc", else flat.
407/// let root = optional_prefix(vb, "esmc", "embed.weight");
408/// # }
409/// ```
410pub fn optional_prefix<'a>(vb: VarBuilder<'a>, prefix: &str, probe: &str) -> VarBuilder<'a> {
411    if vb.contains_tensor(&format!("{prefix}.{probe}")) {
412        vb.pp(prefix)
413    } else {
414        vb
415    }
416}
417
418// ── Tests ─────────────────────────────────────────────────────────────────────
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423
424    #[test]
425    fn test_owner_and_name_splits_valid_id() {
426        let src = WeightSource::safetensors("chandar-lab/AMPLIFY_120M");
427        assert_eq!(
428            src.owner_and_name().unwrap(),
429            ("chandar-lab", "AMPLIFY_120M")
430        );
431    }
432
433    /// The old `split_once('/').unwrap_or(("", repo_id))` turned each of these
434    /// into an empty owner and a confusing 404 after a network round-trip.
435    #[test]
436    fn test_owner_and_name_rejects_malformed_ids() {
437        for bad in ["no-slash", "/leading", "trailing/", "a/b/c", ""] {
438            let err = WeightSource::safetensors(bad)
439                .owner_and_name()
440                .expect_err("{bad} should be rejected");
441            assert!(
442                err.to_string().contains("malformed HuggingFace repo id"),
443                "error should name the malformed id for {bad:?}; got: {err}"
444            );
445        }
446    }
447
448    #[test]
449    fn test_revision_defaults_to_none_and_is_pinnable() {
450        let src = WeightSource::safetensors("owner/name");
451        assert_eq!(src.revision, None);
452        assert_eq!(src.at_revision("v2").revision, Some("v2"));
453    }
454
455    #[test]
456    fn test_format_constructors() {
457        assert_eq!(WeightSource::safetensors("o/n").format, Format::Safetensors);
458        assert_eq!(
459            WeightSource::pth("o/n", Some("model_state_dict")).format,
460            Format::Pth {
461                root_key: Some("model_state_dict")
462            }
463        );
464    }
465
466    /// BF16 on CPU is refused up front with an explanation rather than
467    /// surfacing candle's bare matmul error mid-load (ferritin-100.9).
468    /// The retry budget must exceed hf-hub's own 10-second lock timeout by
469    /// enough to cover a large download, or a concurrent waiter still gives up
470    /// before the downloader finishes (ferritin-100.25).
471    #[test]
472    fn test_cache_lock_retry_budget_outlasts_a_large_download() {
473        assert!(
474            CACHE_LOCK_RETRY_BUDGET >= Duration::from_secs(300),
475            "budget must cover a multi-hundred-MB download; hf-hub's own lock \
476             timeout is only 10s and is not configurable"
477        );
478        assert!(
479            CACHE_LOCK_RETRY_PAUSE < CACHE_LOCK_RETRY_BUDGET,
480            "pause must be shorter than the budget or only one retry happens"
481        );
482    }
483
484    /// A shard index resolves to its unique shard files, sorted and deduped.
485    ///
486    /// The 6-shard ESMC-6B index maps 808 tensors onto 6 files, so the
487    /// many-tensors-to-one-shard collapse is the case that matters.
488    #[test]
489    fn test_shard_names_dedupes_and_sorts() {
490        let dir = std::env::temp_dir().join("ferritin-shard-index-test");
491        std::fs::create_dir_all(&dir).unwrap();
492        let path = dir.join("model.safetensors.index.json");
493        std::fs::write(
494            &path,
495            r#"{
496                "metadata": {"total_size": 123},
497                "weight_map": {
498                    "b.weight": "model-00002-of-00002.safetensors",
499                    "a.weight": "model-00001-of-00002.safetensors",
500                    "a.bias":   "model-00001-of-00002.safetensors"
501                }
502            }"#,
503        )
504        .unwrap();
505
506        let shards = WeightSource::shard_names(&path).unwrap();
507        assert_eq!(
508            shards,
509            [
510                "model-00001-of-00002.safetensors",
511                "model-00002-of-00002.safetensors"
512            ],
513            "three tensors across two shards should yield two unique files"
514        );
515    }
516
517    #[test]
518    fn test_shard_names_rejects_a_non_index() {
519        let dir = std::env::temp_dir().join("ferritin-shard-index-test");
520        std::fs::create_dir_all(&dir).unwrap();
521        let path = dir.join("not-an-index.json");
522        std::fs::write(&path, r#"{"hello": "world"}"#).unwrap();
523
524        let err = WeightSource::shard_names(&path)
525            .map(|_| ())
526            .expect_err("a file without weight_map is not a shard index");
527        assert!(
528            err.to_string().contains("weight_map"),
529            "error should say what was expected; got: {err}"
530        );
531    }
532
533    #[test]
534    fn test_shard_names_rejects_an_empty_weight_map() {
535        let dir = std::env::temp_dir().join("ferritin-shard-index-test");
536        std::fs::create_dir_all(&dir).unwrap();
537        let path = dir.join("empty-index.json");
538        std::fs::write(&path, r#"{"weight_map": {}}"#).unwrap();
539
540        let err = WeightSource::shard_names(&path)
541            .map(|_| ())
542            .expect_err("an empty weight_map names no shards");
543        assert!(err.to_string().contains("empty weight_map"));
544    }
545
546    /// candle resolves tensors across several mmapped shards.
547    ///
548    /// This is the mechanism ESMC-6B needs, exercised on real safetensors
549    /// files rather than assumed. The 6B checkpoint itself cannot be loaded
550    /// here — it is ~24 GB at F32 and ~12 GB at F16 — so this stands in for
551    /// the multi-file part of it (ferritin-100.24).
552    #[test]
553    fn test_var_builder_resolves_tensors_across_shards() {
554        use candle_core::Tensor;
555        use std::collections::HashMap;
556
557        let dir = std::env::temp_dir().join("ferritin-shard-vb-test");
558        std::fs::create_dir_all(&dir).unwrap();
559        let device = Device::Cpu;
560
561        let shard_a = dir.join("model-00001-of-00002.safetensors");
562        let shard_b = dir.join("model-00002-of-00002.safetensors");
563
564        let mut a = HashMap::new();
565        a.insert(
566            "block.0.weight".to_string(),
567            Tensor::zeros((2, 3), DType::F32, &device).unwrap(),
568        );
569        candle_core::safetensors::save(&a, &shard_a).unwrap();
570
571        let mut b = HashMap::new();
572        b.insert(
573            "block.1.weight".to_string(),
574            Tensor::ones((4, 5), DType::F32, &device).unwrap(),
575        );
576        candle_core::safetensors::save(&b, &shard_b).unwrap();
577
578        // SAFETY: files this test just wrote and does not mutate.
579        let vb = unsafe {
580            VarBuilder::from_mmaped_safetensors(&[shard_a, shard_b], DType::F32, &device).unwrap()
581        };
582
583        assert!(
584            vb.contains_tensor("block.0.weight"),
585            "a tensor from the first shard should resolve"
586        );
587        assert!(
588            vb.contains_tensor("block.1.weight"),
589            "a tensor from the second shard should resolve"
590        );
591        assert_eq!(vb.get((2, 3), "block.0.weight").unwrap().dims(), &[2, 3]);
592        assert_eq!(vb.get((4, 5), "block.1.weight").unwrap().dims(), &[4, 5]);
593    }
594
595    /// A `.pth` source cannot be sharded this way, and says so rather than
596    /// producing a confusing mmap failure.
597    #[test]
598    fn test_sharded_loading_is_refused_for_pth() {
599        let err = WeightSource::pth("owner/name", None)
600            .var_builder_sharded(
601                "model.safetensors.index.json",
602                &LoadOptions::new(Device::Cpu),
603            )
604            .map(|_| ())
605            .expect_err("pth sources cannot be sharded");
606        assert!(
607            err.to_string().contains("only defined for safetensors"),
608            "got: {err}"
609        );
610    }
611
612    #[test]
613    fn test_load_options_rejects_bf16_on_cpu() {
614        let err = LoadOptions::new(Device::Cpu)
615            .with_dtype(DType::BF16)
616            .validate()
617            .expect_err("BF16 on CPU must be refused");
618        assert!(
619            err.to_string().contains("not supported on the CPU backend"),
620            "error should explain the CPU limitation; got: {err}"
621        );
622    }
623
624    #[test]
625    fn test_load_options_accepts_f16_and_f32_on_cpu() {
626        for dtype in [DType::F32, DType::F16] {
627            LoadOptions::new(Device::Cpu)
628                .with_dtype(dtype)
629                .validate()
630                .unwrap_or_else(|e| panic!("{dtype:?} should be allowed on CPU: {e}"));
631        }
632    }
633
634    #[test]
635    fn test_load_options_defaults_to_f32() {
636        let opts = LoadOptions::new(Device::Cpu);
637        assert_eq!(opts.dtype, DType::F32);
638        assert_eq!(opts.with_dtype(DType::F16).dtype, DType::F16);
639    }
640
641    /// `optional_prefix` picks the prefixed root only when the probe resolves.
642    #[test]
643    fn test_optional_prefix_falls_back_when_absent() {
644        use candle_core::Tensor;
645        use std::collections::HashMap;
646
647        let device = Device::Cpu;
648        let mut flat = HashMap::new();
649        flat.insert(
650            "embed.weight".to_string(),
651            Tensor::zeros((2, 2), DType::F32, &device).unwrap(),
652        );
653        let vb = VarBuilder::from_tensors(flat, DType::F32, &device);
654        // No "esmc." prefix present → unchanged, so "embed.weight" resolves.
655        let root = optional_prefix(vb, "esmc", "embed.weight");
656        assert!(root.contains_tensor("embed.weight"));
657    }
658
659    #[test]
660    fn test_optional_prefix_descends_when_present() {
661        use candle_core::Tensor;
662        use std::collections::HashMap;
663
664        let device = Device::Cpu;
665        let mut wrapped = HashMap::new();
666        wrapped.insert(
667            "esmc.embed.weight".to_string(),
668            Tensor::zeros((2, 2), DType::F32, &device).unwrap(),
669        );
670        let vb = VarBuilder::from_tensors(wrapped, DType::F32, &device);
671        let root = optional_prefix(vb, "esmc", "embed.weight");
672        assert!(
673            root.contains_tensor("embed.weight"),
674            "after descending into 'esmc' the leaf should resolve unprefixed"
675        );
676    }
677}
ferritin_plms/loader.rs

ferritin_plms/
loader.rs
