tpt_lora_merge/lib.rs
1//! CPU-based merging of LoRA (Low-Rank Adaptation) adapters into base-model
2//! weights, reading and writing the safetensors format.
3//!
4//! The core primitive is [`merge_linear`], which computes
5//! `base + scale * (B @ A)` for a single linear weight. The higher-level
6//! [`merge_lora`] (single adapter) and [`merge_loras`] (weighted sum of several
7//! adapters) walk a base safetensors file and fold in every matching LoRA pair,
8//! producing a [`MergedWeights`] you can serialise back to disk. Merged tensors
9//! preserve the base dtype, base `__metadata__` is carried through, and unused
10//! or partial adapter tensors are reported as errors.
11//!
12//! # Naming convention
13//!
14//! A base weight `"<module>.weight"` of shape `(out, in)` is paired with LoRA
15//! tensors using either the HF PEFT (`".lora_A.weight"`/`".lora_B.weight"`) or
16//! Kohya (`".lora_down.weight"`/`".lora_up.weight"`) convention, plus optional
17//! PEFT named adapters (`".lora_A.<name>.weight"`) and a per-module `".alpha"`
18//! tensor. Any base tensor without a matching LoRA pair is copied through
19//! unchanged.
20//!
21//! # Example
22//!
23//! ```
24//! use ndarray::array;
25//! use tpt_lora_merge::merge_linear;
26//!
27//! let base = array![[1.0_f32, 1.0], [1.0, 1.0]];
28//! let a = array![[1.0_f32, 0.0], [0.0, 1.0]]; // (r, in)
29//! let b = array![[2.0_f32, 0.0], [0.0, 2.0]]; // (out, r)
30//! let merged = merge_linear(base.view(), a.view(), b.view(), 0.5);
31//! assert_eq!(merged, array![[2.0, 1.0], [1.0, 2.0]]);
32//! ```
33#![warn(missing_docs)]
34
35pub mod error;
36pub mod merge;
37
38pub use error::MergeError;
39pub use merge::{merge_linear, merge_lora, merge_loras, MergedWeights};