@ondeinference / cli / commits / 82a845b

Add GGUF export and LoRA merge modules

- Add `gguf.rs` for exporting merged models to GGUF format - Add `merge.rs` for merging LoRA adapters into base models - Update dependencies to include `half` - Update `main.rs` to include new modules - Extend UI and app structures for adapter and GGUF support

keypair34 committed Apr 22, 2026 at 21:35 UTC 82a845b4fc210a012b1e21b622abc1b1fc9599a8
7 files changed +1567 -27
Cargo.lock
+1
@@ -2076,6 +2076,7 @@ dependencies = [
2076 "dirs",
2077 "dotenvy",
2078 "futures",
2079 + "half",
2080 "libc",
2081 "log",
2082 "ratatui",
Cargo.toml
+1
@@ -38,6 +38,7 @@ candle-nn = "0.10.2"
38 tokenizers = { version = "0.20", default-features = false, features = ["onig"] }
39 serde = { version = "1", features = ["derive"] }
40 serde_json = "1"
41 +half = "2"
42
43 # Metal and Accelerate are macOS-only frameworks; Cargo merges features so on
44 # macOS these are added on top of the base candle-core dep above.
src/app.rs
+290 -6
@@ -44,6 +44,19 @@ pub enum Screen {
44 FineTune,
45 }
46
47 +/// A discovered LoRA adapter on disk.
48 +#[derive(Clone)]
49 +pub struct AdapterEntry {
50 + /// Full path to the `lora_adapter.safetensors` file.
51 + pub path: std::path::PathBuf,
52 + /// Directory name (snapshot hash or "lora-adapter").
53 + pub dir_name: String,
54 + /// Human-readable file size.
55 + pub size: String,
56 + /// Relative timestamp ("just now", "3h ago", etc.).
57 + pub modified: String,
58 +}
59 +
60 #[derive(Debug, Clone, PartialEq)]
61 pub enum Mode {
62 Signup,
@@ -142,6 +155,9 @@ pub enum AuthEvent {
155 ModelDownloadFailed(String),
156 // fine-tune
157 FineTuneProgress(crate::finetune::FineTuneProgress),
158 + // merge + GGUF export
159 + MergeProgress(crate::merge::MergeProgress),
160 + GgufProgress(crate::gguf::GgufProgress),
161 }
162
163 pub struct App {
@@ -198,6 +214,15 @@ pub struct App {
214 pub finetune_focus: FineTuneFocus,
215 pub finetune_running: bool,
216 pub finetune_progress: Option<crate::finetune::FineTuneProgress>,
217 + // merge + GGUF export
218 + pub merge_running: bool,
219 + pub merge_progress: Option<crate::merge::MergeProgress>,
220 + pub gguf_running: bool,
221 + pub gguf_progress: Option<crate::gguf::GgufProgress>,
222 + pub merged_model_dir: Option<std::path::PathBuf>,
223 + // adapters discovered on the model detail screen
224 + pub adapter_list: Vec<AdapterEntry>,
225 + pub adapter_cursor: usize,
226 }
227
228 impl App {
@@ -245,6 +270,13 @@ impl App {
270 finetune_focus: FineTuneFocus::ModelDir,
271 finetune_running: false,
272 finetune_progress: None,
273 + merge_running: false,
274 + merge_progress: None,
275 + gguf_running: false,
276 + gguf_progress: None,
277 + merged_model_dir: None,
278 + adapter_list: Vec::new(),
279 + adapter_cursor: 0,
280 }
281 }
282
@@ -266,13 +298,15 @@ impl App {
298 }
299
300 pub fn apply(&mut self, event: AuthEvent) {
269 - // Download and search events run in the background without the busy flag.
301 + // Download, search, merge, and gguf events run in the background without the busy flag.
302 match &event {
303 AuthEvent::ModelDownloadProgress(_)
304 | AuthEvent::ModelDownloadComplete(_)
305 | AuthEvent::ModelDownloadFailed(_)
306 | AuthEvent::HfSearchResults(_)
275 - | AuthEvent::HfSearchFailed(_) => {}
307 + | AuthEvent::HfSearchFailed(_)
308 + | AuthEvent::MergeProgress(_)
309 + | AuthEvent::GgufProgress(_) => {}
310 _ => {
311 self.busy = false;
312 }
@@ -457,6 +491,32 @@ impl App {
491 self.download_progress = None;
492 self.status = Status::error(msg);
493 }
494 + AuthEvent::MergeProgress(progress) => {
495 + match &progress {
496 + crate::merge::MergeProgress::Done { output_path } => {
497 + self.merge_running = false;
498 + // Store the merged model directory for GGUF export
499 + if let Some(parent) = output_path.parent() {
500 + self.merged_model_dir = Some(parent.to_path_buf());
501 + }
502 + }
503 + crate::merge::MergeProgress::Failed(_) => {
504 + self.merge_running = false;
505 + }
506 + _ => {}
507 + }
508 + self.merge_progress = Some(progress);
509 + }
510 + AuthEvent::GgufProgress(progress) => {
511 + match &progress {
512 + crate::gguf::GgufProgress::Done { .. }
513 + | crate::gguf::GgufProgress::Failed(_) => {
514 + self.gguf_running = false;
515 + }
516 + _ => {}
517 + }
518 + self.gguf_progress = Some(progress);
519 + }
520 }
521 }
522 }
@@ -574,7 +634,7 @@ fn handle_key(
634 Screen::AppDetail => handle_key_app_detail(app, key, tx),
635 Screen::Models => handle_key_models(app, key, tx),
636 Screen::Downloads => handle_key_downloads(app, key, tx),
577 - Screen::ModelDetail => handle_key_model_detail(app, key),
637 + Screen::ModelDetail => handle_key_model_detail(app, key, tx),
638 Screen::FineTune => handle_key_finetune(app, key, tx),
639 }
640 }
@@ -906,8 +966,22 @@ fn handle_key_downloads(
966 clamp_downloads_scroll(app, MAX_VISIBLE);
967 }
968 (Enter, _) if !app.downloads.is_empty() => {
969 + // Scan for existing adapters before entering the detail screen
970 + if let Some(model) = app.downloads.get(app.downloads_cursor) {
971 + let resolved = resolve_hf_cache_path(&model.model_id);
972 + app.adapter_list = scan_adapters(&resolved);
973 + app.adapter_cursor = 0;
974 + }
975 app.screen = Screen::ModelDetail;
910 - app.status = Status::neutral("Model details.");
976 + if app.adapter_list.is_empty() {
977 + app.status = Status::neutral("Model details.");
978 + } else {
979 + app.status = Status::success(format!(
980 + "{} adapter{} found. Enter · merge & export f · fine-tune new",
981 + app.adapter_list.len(),
982 + if app.adapter_list.len() == 1 { "" } else { "s" }
983 + ));
984 + }
985 }
986 (Char('/'), KeyModifiers::NONE) => {
987 app.hf_search_active = true;
@@ -928,17 +1002,63 @@ fn handle_key_downloads(
1002 }
1003 }
1004
931 -fn handle_key_model_detail(app: &mut App, key: crossterm::event::KeyEvent) {
1005 +fn handle_key_model_detail(
1006 + app: &mut App,
1007 + key: crossterm::event::KeyEvent,
1008 + _tx: mpsc::UnboundedSender<AuthEvent>,
1009 +) {
1010 use KeyCode::*;
1011
1012 match (key.code, key.modifiers) {
1013 (Esc, _) => {
1014 app.screen = Screen::Downloads;
1015 + app.adapter_list.clear();
1016 + app.adapter_cursor = 0;
1017 app.status = Status::neutral("Back to models.");
1018 }
1019 (Char('c'), KeyModifiers::CONTROL) => {
1020 app.should_quit = true;
1021 }
1022 + // Navigate adapter list
1023 + (Up, _) | (Char('k'), KeyModifiers::NONE) => {
1024 + app.adapter_cursor = app.adapter_cursor.saturating_sub(1);
1025 + }
1026 + (Down, _) | (Char('j'), KeyModifiers::NONE) => {
1027 + if !app.adapter_list.is_empty() && app.adapter_cursor + 1 < app.adapter_list.len() {
1028 + app.adapter_cursor += 1;
1029 + }
1030 + }
1031 + // Enter or 'm' — merge the selected adapter (skip fine-tuning)
1032 + (Enter, _) | (Char('m'), KeyModifiers::NONE) if !app.adapter_list.is_empty() => {
1033 + if let Some(adapter) = app.adapter_list.get(app.adapter_cursor) {
1034 + if let Some(model) = app.downloads.get(app.downloads_cursor) {
1035 + let resolved = resolve_hf_cache_path(&model.model_id);
1036 + if resolved.is_empty() {
1037 + app.status = Status::error("Model not downloaded locally.");
1038 + return;
1039 + }
1040 +
1041 + let adapter_path = adapter.path.clone();
1042 + app.finetune_model_id = model.model_id.clone();
1043 + app.finetune_model_dir = resolved.clone();
1044 +
1045 + // Set finetune_progress to Done so the FineTune screen
1046 + // shows the post-training view with merge/export actions.
1047 + app.finetune_progress = Some(crate::finetune::FineTuneProgress::Done {
1048 + adapter_path: adapter_path.clone(),
1049 + });
1050 + app.finetune_running = false;
1051 + app.merge_progress = None;
1052 + app.gguf_progress = None;
1053 + app.merged_model_dir = None;
1054 + app.screen = Screen::FineTune;
1055 + app.status =
1056 + Status::neutral("Adapter selected. Press m to merge, Esc to go back.");
1057 + }
1058 + } else {
1059 + app.status = Status::error("No adapter selected.");
1060 + }
1061 + }
1062 (Char('f'), KeyModifiers::NONE) => {
1063 if let Some(model) = app.downloads.get(app.downloads_cursor) {
1064 let resolved = resolve_hf_cache_path(&model.model_id);
@@ -1016,6 +1136,86 @@ fn resolve_hf_cache_path(model_id: &str) -> String {
1136 String::new()
1137 }
1138
1139 +/// Scan for existing `lora_adapter.safetensors` files under the model's snapshots
1140 +/// directory. Returns entries sorted by modification time (newest first).
1141 +fn scan_adapters(model_dir: &str) -> Vec<AdapterEntry> {
1142 + if model_dir.is_empty() {
1143 + return Vec::new();
1144 + }
1145 + let snapshots_dir = std::path::Path::new(model_dir)
1146 + .parent()
1147 + .unwrap_or_else(|| std::path::Path::new(""));
1148 +
1149 + let mut adapters: Vec<AdapterEntry> = Vec::new();
1150 +
1151 + let Ok(entries) = std::fs::read_dir(snapshots_dir) else {
1152 + return adapters;
1153 + };
1154 +
1155 + for entry in entries.flatten() {
1156 + let dir = entry.path();
1157 + if !dir.is_dir() {
1158 + continue;
1159 + }
1160 + let adapter_file = dir.join("lora_adapter.safetensors");
1161 + if !adapter_file.exists() {
1162 + continue;
1163 + }
1164 + let dir_name = entry.file_name().to_string_lossy().to_string();
1165 + let meta = std::fs::metadata(&adapter_file);
1166 + let size = meta
1167 + .as_ref()
1168 + .map(|m| format_adapter_size(m.len()))
1169 + .unwrap_or_else(|_| "–".to_string());
1170 + let modified = meta
1171 + .ok()
1172 + .and_then(|m| m.modified().ok())
1173 + .map(|t| {
1174 + let elapsed = t.elapsed().unwrap_or_default();
1175 + let secs = elapsed.as_secs();
1176 + if secs < 60 {
1177 + "just now".to_string()
1178 + } else if secs < 3600 {
1179 + format!("{}m ago", secs / 60)
1180 + } else if secs < 86400 {
1181 + format!("{}h ago", secs / 3600)
1182 + } else {
1183 + format!("{}d ago", secs / 86400)
1184 + }
1185 + })
1186 + .unwrap_or_else(|| "–".to_string());
1187 + adapters.push(AdapterEntry {
1188 + path: adapter_file,
1189 + dir_name,
1190 + size,
1191 + modified,
1192 + });
1193 + }
1194 +
1195 + // Sort newest first (largest mtime string first is unreliable, use the file)
1196 + adapters.sort_by(|a, b| {
1197 + let ma = std::fs::metadata(&a.path)
1198 + .ok()
1199 + .and_then(|m| m.modified().ok());
1200 + let mb = std::fs::metadata(&b.path)
1201 + .ok()
1202 + .and_then(|m| m.modified().ok());
1203 + mb.cmp(&ma)
1204 + });
1205 +
1206 + adapters
1207 +}
1208 +
1209 +fn format_adapter_size(bytes: u64) -> String {
1210 + if bytes >= 1_000_000_000 {
1211 + format!("{:.1}GB", bytes as f64 / 1e9)
1212 + } else if bytes >= 1_000_000 {
1213 + format!("{:.0}MB", bytes as f64 / 1e6)
1214 + } else {
1215 + format!("{:.0}KB", bytes as f64 / 1e3)
1216 + }
1217 +}
1218 +
1219 fn handle_key_finetune(
1220 app: &mut App,
1221 key: crossterm::event::KeyEvent,
@@ -1023,7 +1223,8 @@ fn handle_key_finetune(
1223 ) {
1224 use KeyCode::*;
1225
1026 - if app.finetune_running {
1226 + // While any background task is running, only allow Ctrl+C
1227 + if app.finetune_running || app.merge_running || app.gguf_running {
1228 if matches!(
1229 (key.code, key.modifiers),
1230 (Char('c'), KeyModifiers::CONTROL)
@@ -1033,6 +1234,89 @@ fn handle_key_finetune(
1234 return;
1235 }
1236
1237 + // After fine-tune is done: handle merge (m) and GGUF export (g)
1238 + if let Some(crate::finetune::FineTuneProgress::Done { adapter_path }) = &app.finetune_progress {
1239 + match (key.code, key.modifiers) {
1240 + (Char('m'), KeyModifiers::NONE) => {
1241 + // Start merge: adapter into base model
1242 + let base_dir = std::path::PathBuf::from({
1243 + let raw = &app.finetune_model_dir;
1244 + dirs::home_dir()
1245 + .and_then(|home| {
1246 + raw.strip_prefix("~/")
1247 + .map(|rest| home.join(rest).to_string_lossy().to_string())
1248 + })
1249 + .unwrap_or_else(|| raw.to_string())
1250 + });
1251 + let output_dir = base_dir.parent().unwrap_or(&base_dir).join("merged-model");
1252 +
1253 + let merge_config = crate::merge::MergeConfig {
1254 + base_dir,
1255 + adapter_path: adapter_path.clone(),
1256 + output_dir,
1257 + };
1258 +
1259 + app.merge_running = true;
1260 + app.merge_progress = None;
1261 + app.status = Status::neutral("Merging adapter…");
1262 +
1263 + let (merge_tx, mut merge_rx) = mpsc::unbounded_channel();
1264 + crate::merge::start_merge(merge_config, merge_tx);
1265 +
1266 + let auth_tx = tx.clone();
1267 + tokio::spawn(async move {
1268 + while let Some(progress) = merge_rx.recv().await {
1269 + let _ = auth_tx.send(AuthEvent::MergeProgress(progress));
1270 + }
1271 + });
1272 + return;
1273 + }
1274 + (Char('g'), KeyModifiers::NONE) if app.merged_model_dir.is_some() => {
1275 + // Start GGUF export from merged model
1276 + let model_dir = app.merged_model_dir.clone().unwrap_or_default();
1277 + let output_path = model_dir
1278 + .parent()
1279 + .unwrap_or(&model_dir)
1280 + .join("model-finetuned-q8_0.gguf");
1281 +
1282 + let gguf_config = crate::gguf::GgufConfig {
1283 + model_dir,
1284 + output_path,
1285 + dtype: crate::gguf::GgufDtype::Q8_0,
1286 + };
1287 +
1288 + app.gguf_running = true;
1289 + app.gguf_progress = None;
1290 + app.status = Status::neutral("Exporting GGUF…");
1291 +
1292 + let (gguf_tx, mut gguf_rx) = mpsc::unbounded_channel();
1293 + crate::gguf::start_gguf_export(gguf_config, gguf_tx);
1294 +
1295 + let auth_tx = tx.clone();
1296 + tokio::spawn(async move {
1297 + while let Some(progress) = gguf_rx.recv().await {
1298 + let _ = auth_tx.send(AuthEvent::GgufProgress(progress));
1299 + }
1300 + });
1301 + return;
1302 + }
1303 + (Esc, _) => {
1304 + app.screen = Screen::ModelDetail;
1305 + app.finetune_progress = None;
1306 + app.merge_progress = None;
1307 + app.gguf_progress = None;
1308 + app.merged_model_dir = None;
1309 + app.status = Status::neutral("Back to model.");
1310 + return;
1311 + }
1312 + (Char('c'), KeyModifiers::CONTROL) => {
1313 + app.should_quit = true;
1314 + return;
1315 + }
1316 + _ => return,
1317 + }
1318 + }
1319 +
1320 match (key.code, key.modifiers) {
1321 (Tab, _) => {
1322 app.finetune_focus = match app.finetune_focus {
src/gguf.rs new
+774
@@ -0,0 +1,774 @@
1 +//! GGUF writer module.
2 +//!
3 +//! Converts a merged safetensors model directory (containing `config.json`,
4 +//! `tokenizer.json`, and `model.safetensors`) into a single `.gguf` file
5 +//! compatible with mistral.rs / llama.cpp.
6 +
7 +use std::collections::HashMap;
8 +use std::io::{BufWriter, Write};
9 +use std::path::PathBuf;
10 +
11 +use anyhow::{Context, Result};
12 +use candle_core::Device;
13 +use candle_core::safetensors::MmapedSafetensors;
14 +use half::f16;
15 +
16 +// ---------------------------------------------------------------------------
17 +// Public types
18 +// ---------------------------------------------------------------------------
19 +
20 +/// Target data type for the exported GGUF tensors.
21 +#[allow(dead_code)]
22 +pub enum GgufDtype {
23 + F16,
24 + Q8_0,
25 +}
26 +
27 +/// Configuration for a GGUF export run.
28 +pub struct GgufConfig {
29 + /// Directory containing `config.json`, `tokenizer.json`, and
30 + /// `model.safetensors`.
31 + pub model_dir: PathBuf,
32 + /// Path where the output `.gguf` file will be written.
33 + pub output_path: PathBuf,
34 + /// Target quantisation type for 2-D weight tensors.
35 + pub dtype: GgufDtype,
36 +}
37 +
38 +/// Progress events emitted by the background export thread.
39 +pub enum GgufProgress {
40 + ReadingModel,
41 + WritingTensor {
42 + index: usize,
43 + total: usize,
44 + name: String,
45 + },
46 + Done {
47 + output_path: PathBuf,
48 + size_bytes: u64,
49 + },
50 + Failed(String),
51 +}
52 +
53 +// ---------------------------------------------------------------------------
54 +// Public entry point
55 +// ---------------------------------------------------------------------------
56 +
57 +pub fn start_gguf_export(config: GgufConfig, tx: tokio::sync::mpsc::UnboundedSender<GgufProgress>) {
58 + std::thread::spawn(move || {
59 + if let Err(e) = run_gguf_export(&config, &tx) {
60 + eprintln!("[gguf] error: {e:#}");
61 + let _ = tx.send(GgufProgress::Failed(format!("{e:#}")));
62 + }
63 + });
64 +}
65 +
66 +// ---------------------------------------------------------------------------
67 +// Private — model / tokenizer config structs
68 +// ---------------------------------------------------------------------------
69 +
70 +#[derive(serde::Deserialize)]
71 +struct ModelConfig {
72 + hidden_size: usize,
73 + num_hidden_layers: usize,
74 + num_attention_heads: usize,
75 + num_key_value_heads: Option<usize>,
76 + intermediate_size: usize,
77 + vocab_size: usize,
78 + #[serde(default = "default_rms_norm_eps")]
79 + rms_norm_eps: f64,
80 + #[serde(default = "default_rope_theta")]
81 + rope_theta: f64,
82 + #[serde(default = "default_max_position_embeddings")]
83 + max_position_embeddings: usize,
84 + #[serde(default)]
85 + tie_word_embeddings: bool,
86 +}
87 +
88 +fn default_rms_norm_eps() -> f64 {
89 + 1e-6
90 +}
91 +fn default_rope_theta() -> f64 {
92 + 10000.0
93 +}
94 +fn default_max_position_embeddings() -> usize {
95 + 32768
96 +}
97 +
98 +#[derive(serde::Deserialize)]
99 +struct TokenizerJson {
100 + model: TokenizerModel,
101 + #[serde(default)]
102 + added_tokens: Vec<AddedToken>,
103 +}
104 +
105 +#[derive(serde::Deserialize)]
106 +struct TokenizerModel {
107 + vocab: HashMap<String, u32>,
108 + #[serde(default, deserialize_with = "deserialize_merges")]
109 + merges: Vec<String>,
110 +}
111 +
112 +/// A single merge entry — either `"a b"` (Qwen2.5) or `["a","b"]` (Qwen3).
113 +#[derive(serde::Deserialize)]
114 +#[serde(untagged)]
115 +enum MergeEntry {
116 + Str(String),
117 + Pair(Vec<String>),
118 +}
119 +
120 +/// Deserialize merges from either format into a uniform `Vec<String>` of
121 +/// space-separated pairs (the format GGUF expects).
122 +fn deserialize_merges<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
123 +where
124 + D: serde::Deserializer<'de>,
125 +{
126 + use serde::Deserialize;
127 + let entries: Vec<MergeEntry> = Vec::deserialize(deserializer)?;
128 + Ok(entries
129 + .into_iter()
130 + .map(|e| match e {
131 + MergeEntry::Str(s) => s,
132 + MergeEntry::Pair(parts) => parts.join(" "),
133 + })
134 + .collect())
135 +}
136 +
137 +#[derive(serde::Deserialize)]
138 +struct AddedToken {
139 + id: u32,
140 + content: String,
141 + special: bool,
142 +}
143 +
144 +// ---------------------------------------------------------------------------
145 +// GGUF constants
146 +// ---------------------------------------------------------------------------
147 +
148 +const GGUF_MAGIC: u32 = 0x46475547;
149 +const GGUF_VERSION: u32 = 3;
150 +const ALIGNMENT: u64 = 32;
151 +
152 +// GGUF tensor type IDs
153 +const GGUF_TYPE_F32: u32 = 0;
154 +const GGUF_TYPE_F16: u32 = 1;
155 +const GGUF_TYPE_Q8_0: u32 = 8;
156 +
157 +// GGUF metadata value type IDs
158 +const GGUF_META_U32: u32 = 4;
159 +const GGUF_META_F32: u32 = 6;
160 +const GGUF_META_STRING: u32 = 8;
161 +const GGUF_META_ARRAY: u32 = 9;
162 +
163 +// ---------------------------------------------------------------------------
164 +// Private — GGUF write helpers
165 +// ---------------------------------------------------------------------------
166 +
167 +fn write_u32(w: &mut impl Write, v: u32) -> std::io::Result<()> {
168 + w.write_all(&v.to_le_bytes())
169 +}
170 +
171 +fn write_u64(w: &mut impl Write, v: u64) -> std::io::Result<()> {
172 + w.write_all(&v.to_le_bytes())
173 +}
174 +
175 +fn write_f32(w: &mut impl Write, v: f32) -> std::io::Result<()> {
176 + w.write_all(&v.to_le_bytes())
177 +}
178 +
179 +fn write_gguf_string(w: &mut impl Write, s: &str) -> std::io::Result<()> {
180 + write_u64(w, s.len() as u64)?;
181 + w.write_all(s.as_bytes())
182 +}
183 +
184 +fn write_kv_string(w: &mut impl Write, key: &str, val: &str) -> std::io::Result<()> {
185 + write_gguf_string(w, key)?;
186 + write_u32(w, GGUF_META_STRING)?;
187 + write_gguf_string(w, val)
188 +}
189 +
190 +fn write_kv_u32(w: &mut impl Write, key: &str, val: u32) -> std::io::Result<()> {
191 + write_gguf_string(w, key)?;
192 + write_u32(w, GGUF_META_U32)?;
193 + write_u32(w, val)
194 +}
195 +
196 +fn write_kv_f32(w: &mut impl Write, key: &str, val: f32) -> std::io::Result<()> {
197 + write_gguf_string(w, key)?;
198 + write_u32(w, GGUF_META_F32)?;
199 + write_f32(w, val)
200 +}
201 +
202 +fn write_kv_string_array(w: &mut impl Write, key: &str, vals: &[String]) -> std::io::Result<()> {
203 + write_gguf_string(w, key)?;
204 + write_u32(w, GGUF_META_ARRAY)?;
205 + write_u32(w, GGUF_META_STRING)?;
206 + write_u64(w, vals.len() as u64)?;
207 + for s in vals {
208 + write_gguf_string(w, s)?;
209 + }
210 + Ok(())
211 +}
212 +
213 +fn write_kv_u32_array(w: &mut impl Write, key: &str, vals: &[u32]) -> std::io::Result<()> {
214 + write_gguf_string(w, key)?;
215 + write_u32(w, GGUF_META_ARRAY)?;
216 + write_u32(w, GGUF_META_U32)?;
217 + write_u64(w, vals.len() as u64)?;
218 + for &v in vals {
219 + write_u32(w, v)?;
220 + }
221 + Ok(())
222 +}
223 +
224 +fn pad_to_alignment(w: &mut impl Write, pos: u64, alignment: u64) -> std::io::Result<u64> {
225 + let rem = pos % alignment;
226 + if rem != 0 {
227 + let pad = alignment - rem;
228 + for _ in 0..pad {
229 + w.write_all(&[0u8])?;
230 + }
231 + Ok(pos + pad)
232 + } else {
233 + Ok(pos)
234 + }
235 +}
236 +
237 +// ---------------------------------------------------------------------------
238 +// Private — tensor name mapping (HuggingFace → GGUF)
239 +// ---------------------------------------------------------------------------
240 +
241 +fn map_tensor_name(hf_name: &str) -> Option<String> {
242 + // Non-layer names
243 + if hf_name == "model.embed_tokens.weight" {
244 + return Some("token_embd.weight".to_string());
245 + }
246 + if hf_name == "model.norm.weight" {
247 + return Some("output_norm.weight".to_string());
248 + }
249 + if hf_name == "lm_head.weight" {
250 + return Some("output.weight".to_string());
251 + }
252 +
253 + // Layer names: model.layers.{i}.suffix → blk.{i}.gguf_suffix
254 + let prefix = "model.layers.";
255 + if !hf_name.starts_with(prefix) {
256 + return None;
257 + }
258 +
259 + let rest = &hf_name[prefix.len()..];
260 + let dot_pos = rest.find('.')?;
261 + let layer_idx = &rest[..dot_pos];
262 + let suffix = &rest[dot_pos + 1..];
263 +
264 + let gguf_suffix = match suffix {
265 + "self_attn.q_proj.weight" => "attn_q.weight",
266 + "self_attn.k_proj.weight" => "attn_k.weight",
267 + "self_attn.v_proj.weight" => "attn_v.weight",
268 + "self_attn.o_proj.weight" => "attn_output.weight",
269 + "self_attn.q_norm.weight" => "attn_q_norm.weight",
270 + "self_attn.k_norm.weight" => "attn_k_norm.weight",
271 + "mlp.gate_proj.weight" => "ffn_gate.weight",
272 + "mlp.up_proj.weight" => "ffn_up.weight",
273 + "mlp.down_proj.weight" => "ffn_down.weight",
274 + "input_layernorm.weight" => "attn_norm.weight",
275 + "post_attention_layernorm.weight" => "ffn_norm.weight",
276 + _ => return None,
277 + };
278 +
279 + Some(format!("blk.{layer_idx}.{gguf_suffix}"))
280 +}
281 +
282 +// ---------------------------------------------------------------------------
283 +// Private — tensor data size calculation
284 +// ---------------------------------------------------------------------------
285 +
286 +fn tensor_data_size(num_elements: usize, gguf_type: u32) -> usize {
287 + match gguf_type {
288 + GGUF_TYPE_F32 => num_elements * 4,
289 + GGUF_TYPE_F16 => num_elements * 2,
290 + GGUF_TYPE_Q8_0 => num_elements.div_ceil(32) * 34,
291 + _ => num_elements * 4,
292 + }
293 +}
294 +
295 +// ---------------------------------------------------------------------------
296 +// Private — determine GGUF type for a tensor
297 +// ---------------------------------------------------------------------------
298 +
299 +fn choose_gguf_type(gguf_name: &str, ndim: usize, dtype: &GgufDtype) -> u32 {
300 + // 1-D tensors (norms, biases) → always F32
301 + if ndim <= 1 {
302 + return GGUF_TYPE_F32;
303 + }
304 +
305 + // Embedding and output head → always F16
306 + if gguf_name == "token_embd.weight" || gguf_name == "output.weight" {
307 + return GGUF_TYPE_F16;
308 + }
309 +
310 + // 2-D weight tensors → follow requested dtype
311 + match dtype {
312 + GgufDtype::F16 => GGUF_TYPE_F16,
313 + GgufDtype::Q8_0 => GGUF_TYPE_Q8_0,
314 + }
315 +}
316 +
317 +// ---------------------------------------------------------------------------
318 +// Private — tensor data conversion and writing
319 +// ---------------------------------------------------------------------------
320 +
321 +fn write_tensor_f32(w: &mut impl Write, data: &[f32]) -> std::io::Result<()> {
322 + for &v in data {
323 + w.write_all(&v.to_le_bytes())?;
324 + }
325 + Ok(())
326 +}
327 +
328 +fn write_tensor_f16(w: &mut impl Write, data: &[f32]) -> std::io::Result<()> {
329 + for &v in data {
330 + w.write_all(&f16::from_f32(v).to_le_bytes())?;
331 + }
332 + Ok(())
333 +}
334 +
335 +fn write_tensor_q8_0(w: &mut impl Write, data: &[f32]) -> std::io::Result<()> {
336 + let block_size = 32;
337 + let num_blocks = data.len().div_ceil(block_size);
338 +
339 + for block_idx in 0..num_blocks {
340 + let start = block_idx * block_size;
341 + let end = (start + block_size).min(data.len());
342 + let block = &data[start..end];
343 +
344 + // Find max absolute value in the block
345 + let mut amax: f32 = 0.0;
346 + for &v in block {
347 + let abs = v.abs();
348 + if abs > amax {
349 + amax = abs;
350 + }
351 + }
352 +
353 + let scale = if amax == 0.0 { 0.0 } else { amax / 127.0 };
354 + let inv_scale = if scale == 0.0 { 0.0 } else { 1.0 / scale };
355 +
356 + // Write scale as f16 (2 bytes)
357 + w.write_all(&f16::from_f32(scale).to_le_bytes())?;
358 +
359 + // Quantize and write 32 i8 values
360 + for i in 0..block_size {
361 + let val = if i < block.len() { block[i] } else { 0.0 };
362 + let q = (val * inv_scale).round().clamp(-128.0, 127.0) as i8;
363 + w.write_all(&[q as u8])?;
364 + }
365 + }
366 + Ok(())
367 +}
368 +
369 +// ---------------------------------------------------------------------------
370 +// Private — compute byte sizes of metadata
371 +// ---------------------------------------------------------------------------
372 +
373 +/// Compute the byte size of a GGUF string (length prefix + bytes).
374 +fn gguf_string_size(s: &str) -> u64 {
375 + 8 + s.len() as u64
376 +}
377 +
378 +/// Compute the byte size of a KV string entry.
379 +fn kv_string_size(key: &str, val: &str) -> u64 {
380 + gguf_string_size(key) + 4 + gguf_string_size(val)
381 +}
382 +
383 +/// Compute the byte size of a KV u32 entry.
384 +fn kv_u32_size(key: &str) -> u64 {
385 + gguf_string_size(key) + 4 + 4
386 +}
387 +
388 +/// Compute the byte size of a KV f32 entry.
389 +fn kv_f32_size(key: &str) -> u64 {
390 + gguf_string_size(key) + 4 + 4
391 +}
392 +
393 +/// Compute the byte size of a KV string-array entry.
394 +fn kv_string_array_size(key: &str, vals: &[String]) -> u64 {
395 + let mut size = gguf_string_size(key) + 4; // key + value_type
396 + size += 4 + 8; // array element_type + count
397 + for s in vals {
398 + size += gguf_string_size(s);
399 + }
400 + size
401 +}
402 +
403 +/// Compute the byte size of a KV u32-array entry.
404 +fn kv_u32_array_size(key: &str, vals: &[u32]) -> u64 {
405 + gguf_string_size(key) + 4 + 4 + 8 + (vals.len() as u64 * 4)
406 +}
407 +
408 +/// Compute the byte size of a single tensor info entry.
409 +fn tensor_info_size(name: &str, ndim: usize) -> u64 {
410 + // name_len(8) + name_bytes + n_dims(4) + dims(8*ndim) + type(4) + offset(8)
411 + 8 + name.len() as u64 + 4 + (ndim as u64 * 8) + 4 + 8
412 +}
413 +
414 +// ---------------------------------------------------------------------------
415 +// Private — core export logic
416 +// ---------------------------------------------------------------------------
417 +
418 +/// Holds the info we need for each tensor in the GGUF file.
419 +struct TensorEntry {
420 + /// GGUF tensor name.
421 + gguf_name: String,
422 + /// Original HuggingFace tensor name (for loading from safetensors).
423 + hf_name: String,
424 + /// Dimensions in GGUF order (row-major, i.e. same order as safetensors).
425 + dims: Vec<u64>,
426 + /// Number of dimensions.
427 + ndim: usize,
428 + /// Total number of elements.
429 + num_elements: usize,
430 + /// GGUF tensor type ID.
431 + gguf_type: u32,
432 + /// Size in bytes of the tensor data.
433 + data_size: usize,
434 + /// Offset from start of data section.
435 + offset: u64,
436 +}
437 +
438 +fn run_gguf_export(
439 + config: &GgufConfig,
440 + tx: &tokio::sync::mpsc::UnboundedSender<GgufProgress>,
441 +) -> Result<()> {
442 + let _ = tx.send(GgufProgress::ReadingModel);
443 +
444 + // -----------------------------------------------------------------------
445 + // 1. Load model config
446 + // -----------------------------------------------------------------------
447 + let config_path = config.model_dir.join("config.json");
448 + let config_text = std::fs::read_to_string(&config_path)
449 + .with_context(|| format!("failed to read {}", config_path.display()))?;
450 + let model_cfg: ModelConfig =
451 + serde_json::from_str(&config_text).context("failed to parse config.json")?;
452 +
453 + // -----------------------------------------------------------------------
454 + // 2. Load and parse tokenizer.json
455 + // -----------------------------------------------------------------------
456 + let tok_path = config.model_dir.join("tokenizer.json");
457 + let tok_text = std::fs::read_to_string(&tok_path)
458 + .with_context(|| format!("failed to read {}", tok_path.display()))?;
459 + let tok: TokenizerJson =
460 + serde_json::from_str(&tok_text).context("failed to parse tokenizer.json")?;
461 +
462 + // Build ordered token list sorted by ID
463 + let vocab_size = model_cfg.vocab_size;
464 + let mut id_to_token: Vec<(u32, String)> =
465 + tok.model.vocab.into_iter().map(|(k, v)| (v, k)).collect();
466 + // Also include added_tokens that might not be in the base vocab
467 + for at in &tok.added_tokens {
468 + if !id_to_token.iter().any(|(id, _)| *id == at.id) {
469 + id_to_token.push((at.id, at.content.clone()));
470 + }
471 + }
472 + id_to_token.sort_by_key(|(id, _)| *id);
473 +
474 + // Pad to vocab_size if needed
475 + let mut tokens: Vec<String> = Vec::with_capacity(vocab_size);
476 + let mut next_expected: u32 = 0;
477 + for (id, content) in &id_to_token {
478 + while next_expected < *id && (next_expected as usize) < vocab_size {
479 + tokens.push(format!("<unused{next_expected}>"));
480 + next_expected += 1;
481 + }
482 + if (next_expected as usize) < vocab_size {
483 + tokens.push(content.clone());
484 + next_expected += 1;
485 + }
486 + }
487 + while tokens.len() < vocab_size {
488 + let idx = tokens.len();
489 + tokens.push(format!("<unused{idx}>"));
490 + }
491 + tokens.truncate(vocab_size);
492 +
493 + // Build special token lookup
494 + let special_set: HashMap<u32, bool> = tok
495 + .added_tokens
496 + .iter()
497 + .filter(|at| at.special)
498 + .map(|at| (at.id, true))
499 + .collect();
500 +
501 + // Token types: 1 = normal, 3 = control (special)
502 + let token_types: Vec<u32> = (0..vocab_size as u32)
503 + .map(|id| if special_set.contains_key(&id) { 3 } else { 1 })
504 + .collect();
505 +
506 + // BOS / EOS token IDs
507 + let added_token_id = |name: &str| -> Option<u32> {
508 + tok.added_tokens
509 + .iter()
510 + .find(|at| at.content == name)
511 + .map(|at| at.id)
512 + };
513 + let bos_id = added_token_id("<|endoftext|>").unwrap_or(0);
514 + let eos_id = added_token_id("<|im_end|>")
515 + .or_else(|| added_token_id("<|endoftext|>"))
516 + .unwrap_or(0);
517 +
518 + let merges = tok.model.merges;
519 +
520 + // -----------------------------------------------------------------------
521 + // 3. Load safetensors
522 + // -----------------------------------------------------------------------
523 + let st_path = config.model_dir.join("model.safetensors");
524 + // SAFETY: the file is memory-mapped read-only.
525 + let safetensors = unsafe { MmapedSafetensors::new(&st_path) }
526 + .with_context(|| format!("failed to mmap {}", st_path.display()))?;
527 + let tensor_list = safetensors.tensors();
528 +
529 + // -----------------------------------------------------------------------
530 + // 4. Build tensor entries with name mapping
531 + // -----------------------------------------------------------------------
532 + let mut entries: Vec<TensorEntry> = Vec::new();
533 + for (hf_name, tensor_view) in &tensor_list {
534 + let gguf_name = match map_tensor_name(hf_name) {
535 + Some(n) => n,
536 + None => continue,
537 + };
538 +
539 + let shape = tensor_view.shape();
540 + let ndim = shape.len();
541 + let num_elements: usize = shape.iter().product();
542 + let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
543 + let gguf_type = choose_gguf_type(&gguf_name, ndim, &config.dtype);
544 + let data_size = tensor_data_size(num_elements, gguf_type);
545 +
546 + entries.push(TensorEntry {
547 + gguf_name,
548 + hf_name: hf_name.clone(),
549 + dims,
550 + ndim,
551 + num_elements,
552 + gguf_type,
553 + data_size,
554 + offset: 0, // computed below
555 + });
556 + }
557 +
558 + // Handle tie_word_embeddings: if lm_head.weight is missing, duplicate
559 + // embed_tokens as output.weight
560 + let has_output = entries.iter().any(|e| e.gguf_name == "output.weight");
561 + if !has_output
562 + && model_cfg.tie_word_embeddings
563 + && let Some(embd) = entries.iter().find(|e| e.gguf_name == "token_embd.weight")
564 + {
565 + let gguf_type = choose_gguf_type("output.weight", embd.ndim, &config.dtype);
566 + let data_size = tensor_data_size(embd.num_elements, gguf_type);
567 + entries.push(TensorEntry {
568 + gguf_name: "output.weight".to_string(),
569 + hf_name: "model.embed_tokens.weight".to_string(),
570 + dims: embd.dims.clone(),
571 + ndim: embd.ndim,
572 + num_elements: embd.num_elements,
573 + gguf_type,
574 + data_size,
575 + offset: 0,
576 + });
577 + }
578 +
579 + // Compute offsets (relative to start of data section, 32-byte aligned)
580 + let mut data_offset: u64 = 0;
581 + for entry in &mut entries {
582 + entry.offset = data_offset;
583 + let size = entry.data_size as u64;
584 + data_offset += size;
585 + // Pad to 32-byte alignment
586 + let rem = data_offset % ALIGNMENT;
587 + if rem != 0 {
588 + data_offset += ALIGNMENT - rem;
589 + }
590 + }
591 +
592 + // -----------------------------------------------------------------------
593 + // 5. Compute header size to know where data section starts
594 + // -----------------------------------------------------------------------
595 + let num_tensors = entries.len() as u64;
596 + let num_kv_heads = model_cfg
597 + .num_key_value_heads
598 + .unwrap_or(model_cfg.num_attention_heads) as u32;
599 + let file_type_val: u32 = match config.dtype {
600 + GgufDtype::F16 => 1,
601 + GgufDtype::Q8_0 => 7,
602 + };
603 +
604 + // Count metadata KV entries
605 + let num_metadata: u64 = 17; // all the keys listed below
606 +
607 + // Pre-compute header + metadata + tensor info byte size
608 + let mut header_size: u64 = 0;
609 + // Magic + version + tensor_count + metadata_kv_count
610 + header_size += 4 + 4 + 8 + 8;
611 +
612 + // Metadata sizes
613 + header_size += kv_string_size("general.architecture", "qwen2");
614 + header_size += kv_string_size("general.name", "onde-finetuned");
615 + header_size += kv_u32_size("general.file_type");
616 + header_size += kv_u32_size("qwen2.block_count");
617 + header_size += kv_u32_size("qwen2.embedding_length");
618 + header_size += kv_u32_size("qwen2.feed_forward_length");
619 + header_size += kv_u32_size("qwen2.attention.head_count");
620 + header_size += kv_u32_size("qwen2.attention.head_count_kv");
621 + header_size += kv_f32_size("qwen2.attention.layer_norm_rms_epsilon");
622 + header_size += kv_f32_size("qwen2.rope.freq_base");
623 + header_size += kv_u32_size("qwen2.context_length");
624 + header_size += kv_string_size("tokenizer.ggml.model", "gpt2");
625 + header_size += kv_string_array_size("tokenizer.ggml.tokens", &tokens);
626 + header_size += kv_u32_array_size("tokenizer.ggml.token_type", &token_types);
627 + header_size += kv_string_array_size("tokenizer.ggml.merges", &merges);
628 + header_size += kv_u32_size("tokenizer.ggml.bos_token_id");
629 + header_size += kv_u32_size("tokenizer.ggml.eos_token_id");
630 +
631 + // Tensor info sizes
632 + for entry in &entries {
633 + header_size += tensor_info_size(&entry.gguf_name, entry.ndim);
634 + }
635 +
636 + // Pad header to alignment
637 + let header_padded = {
638 + let rem = header_size % ALIGNMENT;
639 + if rem != 0 {
640 + header_size + (ALIGNMENT - rem)
641 + } else {
642 + header_size
643 + }
644 + };
645 +
646 + // -----------------------------------------------------------------------
647 + // 6. Write the GGUF file
648 + // -----------------------------------------------------------------------
649 + let file = std::fs::File::create(&config.output_path)
650 + .with_context(|| format!("failed to create {}", config.output_path.display()))?;
651 + let mut w = BufWriter::new(file);
652 +
653 + // Header
654 + write_u32(&mut w, GGUF_MAGIC)?;
655 + write_u32(&mut w, GGUF_VERSION)?;
656 + write_u64(&mut w, num_tensors)?;
657 + write_u64(&mut w, num_metadata)?;
658 +
659 + // Metadata KV pairs
660 + write_kv_string(&mut w, "general.architecture", "qwen2")?;
661 + write_kv_string(&mut w, "general.name", "onde-finetuned")?;
662 + write_kv_u32(&mut w, "general.file_type", file_type_val)?;
663 + write_kv_u32(
664 + &mut w,
665 + "qwen2.block_count",
666 + model_cfg.num_hidden_layers as u32,
667 + )?;
668 + write_kv_u32(
669 + &mut w,
670 + "qwen2.embedding_length",
671 + model_cfg.hidden_size as u32,
672 + )?;
673 + write_kv_u32(
674 + &mut w,
675 + "qwen2.feed_forward_length",
676 + model_cfg.intermediate_size as u32,
677 + )?;
678 + write_kv_u32(
679 + &mut w,
680 + "qwen2.attention.head_count",
681 + model_cfg.num_attention_heads as u32,
682 + )?;
683 + write_kv_u32(&mut w, "qwen2.attention.head_count_kv", num_kv_heads)?;
684 + write_kv_f32(
685 + &mut w,
686 + "qwen2.attention.layer_norm_rms_epsilon",
687 + model_cfg.rms_norm_eps as f32,
688 + )?;
689 + write_kv_f32(&mut w, "qwen2.rope.freq_base", model_cfg.rope_theta as f32)?;
690 + write_kv_u32(
691 + &mut w,
692 + "qwen2.context_length",
693 + model_cfg.max_position_embeddings as u32,
694 + )?;
695 + write_kv_string(&mut w, "tokenizer.ggml.model", "gpt2")?;
696 + write_kv_string_array(&mut w, "tokenizer.ggml.tokens", &tokens)?;
697 + write_kv_u32_array(&mut w, "tokenizer.ggml.token_type", &token_types)?;
698 + write_kv_string_array(&mut w, "tokenizer.ggml.merges", &merges)?;
699 + write_kv_u32(&mut w, "tokenizer.ggml.bos_token_id", bos_id)?;
700 + write_kv_u32(&mut w, "tokenizer.ggml.eos_token_id", eos_id)?;
701 +
702 + // Tensor info entries
703 + for entry in &entries {
704 + write_gguf_string(&mut w, &entry.gguf_name)?;
705 + write_u32(&mut w, entry.ndim as u32)?;
706 + for &d in &entry.dims {
707 + write_u64(&mut w, d)?;
708 + }
709 + write_u32(&mut w, entry.gguf_type)?;
710 + write_u64(&mut w, entry.offset)?;
711 + }
712 +
713 + // Pad to alignment before data section
714 + let mut pos = header_size;
715 + pos = pad_to_alignment(&mut w, pos, ALIGNMENT)?;
716 + let _ = header_padded; // should match pos
717 +
718 + // -----------------------------------------------------------------------
719 + // 7. Write tensor data
720 + // -----------------------------------------------------------------------
721 + let total = entries.len();
722 + for (idx, entry) in entries.iter().enumerate() {
723 + let _ = tx.send(GgufProgress::WritingTensor {
724 + index: idx + 1,
725 + total,
726 + name: entry.gguf_name.clone(),
727 + });
728 +
729 + // Load tensor from safetensors and convert to f32 on CPU
730 + let tensor = safetensors
731 + .load(&entry.hf_name, &Device::Cpu)
732 + .with_context(|| format!("failed to load tensor {}", entry.hf_name))?;
733 + let tensor_f32 = tensor
734 + .to_dtype(candle_core::DType::F32)
735 + .with_context(|| format!("failed to convert {} to f32", entry.hf_name))?;
736 + let data = tensor_f32
737 + .flatten_all()
738 + .context("failed to flatten tensor")?;
739 + let data = data
740 + .to_vec1::<f32>()
741 + .context("failed to extract f32 data from tensor")?;
742 +
743 + match entry.gguf_type {
744 + GGUF_TYPE_F32 => write_tensor_f32(&mut w, &data)?,
745 + GGUF_TYPE_F16 => write_tensor_f16(&mut w, &data)?,
746 + GGUF_TYPE_Q8_0 => write_tensor_q8_0(&mut w, &data)?,
747 + _ => write_tensor_f32(&mut w, &data)?,
748 + }
749 +
750 + pos += entry.data_size as u64;
751 +
752 + // Pad to alignment (except possibly the last tensor)
753 + if idx + 1 < total {
754 + pos = pad_to_alignment(&mut w, pos, ALIGNMENT)?;
755 + }
756 + }
757 +
758 + w.flush().context("failed to flush GGUF output")?;
759 + drop(w);
760 +
761 + // -----------------------------------------------------------------------
762 + // 8. Report completion
763 + // -----------------------------------------------------------------------
764 + let size_bytes = std::fs::metadata(&config.output_path)
765 + .with_context(|| format!("failed to stat {}", config.output_path.display()))?
766 + .len();
767 +
768 + let _ = tx.send(GgufProgress::Done {
769 + output_path: config.output_path.clone(),
770 + size_bytes,
771 + });
772 +
773 + Ok(())
774 +}
src/main.rs
+2
@@ -1,8 +1,10 @@
1 mod app;
2 mod finetune;
3 +mod gguf;
4 mod gresiq;
5 mod hf;
6 mod hf_search;
7 +mod merge;
8 mod token;
9 mod ui;
10
src/merge.rs new
+233
@@ -0,0 +1,233 @@
1 +//! LoRA adapter merge module.
2 +//!
3 +//! Merges LoRA adapter weights (produced by the fine-tuner) back into the base
4 +//! model safetensors, producing a single merged `model.safetensors` alongside
5 +//! copied tokenizer and config files.
6 +
7 +use std::collections::HashMap;
8 +use std::path::PathBuf;
9 +
10 +use anyhow::{Context, Result};
11 +use candle_core::safetensors::MmapedSafetensors;
12 +use candle_core::{DType, Device, Tensor};
13 +
14 +// ---------------------------------------------------------------------------
15 +// Public types
16 +// ---------------------------------------------------------------------------
17 +
18 +pub struct MergeConfig {
19 + /// Snapshot directory containing `model.safetensors` (or sharded files) and
20 + /// `config.json`.
21 + pub base_dir: PathBuf,
22 + /// Path to the `lora_adapter.safetensors` produced by the fine-tuner.
23 + pub adapter_path: PathBuf,
24 + /// Directory where the merged model and copied config/tokenizer files will
25 + /// be written.
26 + pub output_dir: PathBuf,
27 +}
28 +
29 +pub enum MergeProgress {
30 + Loading,
31 + Merging { layer: usize, total: usize },
32 + Saving,
33 + Done { output_path: PathBuf },
34 + Failed(String),
35 +}
36 +
37 +// ---------------------------------------------------------------------------
38 +// Public entry point
39 +// ---------------------------------------------------------------------------
40 +
41 +pub fn start_merge(config: MergeConfig, tx: tokio::sync::mpsc::UnboundedSender<MergeProgress>) {
42 + std::thread::spawn(move || {
43 + if let Err(e) = run_merge(&config, &tx) {
44 + eprintln!("[merge] error: {e:#}");
45 + let _ = tx.send(MergeProgress::Failed(format!("{e:#}")));
46 + }
47 + });
48 +}
49 +
50 +// ---------------------------------------------------------------------------
51 +// Private helpers
52 +// ---------------------------------------------------------------------------
53 +
54 +#[derive(serde::Deserialize)]
55 +struct MinimalConfig {
56 + num_hidden_layers: usize,
57 +}
58 +
59 +#[derive(serde::Deserialize)]
60 +struct IndexJson {
61 + weight_map: HashMap<String, String>,
62 +}
63 +
64 +fn run_merge(
65 + config: &MergeConfig,
66 + tx: &tokio::sync::mpsc::UnboundedSender<MergeProgress>,
67 +) -> Result<()> {
68 + let _ = tx.send(MergeProgress::Loading);
69 +
70 + // 1. Read num_hidden_layers from config.json
71 + let config_path = config.base_dir.join("config.json");
72 + let config_text = std::fs::read_to_string(&config_path)
73 + .with_context(|| format!("failed to read {}", config_path.display()))?;
74 + let model_cfg: MinimalConfig =
75 + serde_json::from_str(&config_text).context("failed to parse config.json")?;
76 + let num_layers = model_cfg.num_hidden_layers;
77 +
78 + // 2. Load base safetensors (single file or sharded)
79 + let index_path = config.base_dir.join("model.safetensors.index.json");
80 + let base_paths: Vec<PathBuf> = if index_path.exists() {
81 + let index_text = std::fs::read_to_string(&index_path)
82 + .with_context(|| format!("failed to read {}", index_path.display()))?;
83 + let index: IndexJson = serde_json::from_str(&index_text)
84 + .context("failed to parse model.safetensors.index.json")?;
85 + let mut filenames: Vec<String> = index.weight_map.into_values().collect();
86 + filenames.sort();
87 + filenames.dedup();
88 + filenames
89 + .into_iter()
90 + .map(|f| config.base_dir.join(f))
91 + .collect()
92 + } else {
93 + vec![config.base_dir.join("model.safetensors")]
94 + };
95 +
96 + // SAFETY: the files are memory-mapped read-only.
97 + let base = if base_paths.len() == 1 {
98 + unsafe { MmapedSafetensors::new(&base_paths[0]) }
99 + } else {
100 + unsafe { MmapedSafetensors::multi(&base_paths) }
101 + }
102 + .context("failed to mmap base safetensors")?;
103 +
104 + // 3. Load adapter tensors
105 + let adapter =
106 + candle_core::safetensors::load(&config.adapter_path, &Device::Cpu).with_context(|| {
107 + format!(
108 + "failed to load adapter from {}",
109 + config.adapter_path.display()
110 + )
111 + })?;
112 +
113 + // 4. Infer rank from any lora_a tensor
114 + let rank = adapter
115 + .iter()
116 + .find(|(k, _)| k.ends_with("lora_a"))
117 + .map(|(_, t)| t.dim(0))
118 + .context("no lora_a tensor found in adapter")?
119 + .context("failed to read lora_a dimension")?;
120 +
121 + // 5. Scale = alpha / rank = (rank * 2) / rank = 2.0
122 + let scale = 2.0_f64;
123 +
124 + // 6. Iterate base tensors and merge where adapters exist
125 + let base_tensor_info = base.tensors();
126 + let mut merged_tensors: HashMap<String, Tensor> = HashMap::new();
127 + let mut merge_layer_idx: usize = 0;
128 +
129 + // Count how many layers actually have adapters for progress reporting
130 + let total_merge_layers = {
131 + let mut count = 0usize;
132 + for i in 0..num_layers {
133 + let q_a = format!("model.layers.{i}.self_attn.q_proj.lora_a");
134 + let v_a = format!("model.layers.{i}.self_attn.v_proj.lora_a");
135 + if adapter.contains_key(&q_a) || adapter.contains_key(&v_a) {
136 + count += 1;
137 + }
138 + }
139 + count
140 + };
141 +
142 + // Track which layers we've already sent progress for
143 + let mut reported_layers: std::collections::HashSet<usize> = std::collections::HashSet::new();
144 +
145 + for (name, _tv) in &base_tensor_info {
146 + let base_tensor = base
147 + .load(name, &Device::Cpu)
148 + .with_context(|| format!("failed to load base tensor {name}"))?;
149 + let base_dtype = base_tensor.dtype();
150 +
151 + let stem = name.strip_suffix(".weight").unwrap_or(name);
152 + let lora_a_key = format!("{stem}.lora_a");
153 + let lora_b_key = format!("{stem}.lora_b");
154 +
155 + let tensor = if let (Some(lora_a), Some(lora_b)) =
156 + (adapter.get(&lora_a_key), adapter.get(&lora_b_key))
157 + {
158 + // Report progress per-layer
159 + // Extract layer index from the stem if possible
160 + if let Some(layer_num) = extract_layer_index(stem)
161 + && reported_layers.insert(layer_num)
162 + {
163 + merge_layer_idx += 1;
164 + let _ = tx.send(MergeProgress::Merging {
165 + layer: merge_layer_idx,
166 + total: total_merge_layers,
167 + });
168 + }
169 +
170 + let base_f32 = base_tensor.to_dtype(DType::F32)?;
171 + let a_f32 = lora_a.to_dtype(DType::F32)?;
172 + let b_f32 = lora_b.to_dtype(DType::F32)?;
173 + let delta = (b_f32.matmul(&a_f32)? * scale)?;
174 + let merged = (base_f32 + delta)?;
175 + merged.to_dtype(base_dtype)?
176 + } else {
177 + base_tensor
178 + };
179 +
180 + merged_tensors.insert(name.clone(), tensor);
181 + }
182 +
183 + // Ensure rank is used (suppress unused warning via a debug log)
184 + let _ = rank;
185 +
186 + // 7. Save merged model
187 + let _ = tx.send(MergeProgress::Saving);
188 + std::fs::create_dir_all(&config.output_dir).with_context(|| {
189 + format!(
190 + "failed to create output dir {}",
191 + config.output_dir.display()
192 + )
193 + })?;
194 +
195 + let output_model = config.output_dir.join("model.safetensors");
196 + candle_core::safetensors::save(&merged_tensors, &output_model)
197 + .with_context(|| format!("failed to save merged model to {}", output_model.display()))?;
198 +
199 + // 8. Copy config/tokenizer files (skip missing)
200 + let copy_files = [
201 + "config.json",
202 + "tokenizer.json",
203 + "tokenizer_config.json",
204 + "generation_config.json",
205 + ];
206 + for filename in &copy_files {
207 + let src = config.base_dir.join(filename);
208 + if src.exists() {
209 + let dst = config.output_dir.join(filename);
210 + if let Err(e) = std::fs::copy(&src, &dst) {
211 + eprintln!("[merge] warning: failed to copy {}: {e:#}", src.display());
212 + }
213 + }
214 + }
215 +
216 + // 9. Done
217 + let _ = tx.send(MergeProgress::Done {
218 + output_path: output_model,
219 + });
220 + Ok(())
221 +}
222 +
223 +/// Extract a numeric layer index from a tensor name stem such as
224 +/// `model.layers.5.self_attn.q_proj`.
225 +fn extract_layer_index(stem: &str) -> Option<usize> {
226 + let parts: Vec<&str> = stem.split('.').collect();
227 + for (i, part) in parts.iter().enumerate() {
228 + if *part == "layers" {
229 + return parts.get(i + 1).and_then(|s| s.parse::<usize>().ok());
230 + }
231 + }
232 + None
233 +}
src/ui.rs
+266 -21
@@ -1041,8 +1041,10 @@ fn render_model_detail(frame: &mut Frame, app: &App, area: Rect) {
1041 Constraint::Length(1), // 23 Catalog ID value
1042 Constraint::Length(1), // 24 spacer
1043 Constraint::Length(1), // 25 Description label
1044 - Constraint::Min(0), // 26 Description value (wraps)
1045 - Constraint::Length(1), // 27 hint
1044 + Constraint::Length(1), // 26 Description value
1045 + Constraint::Length(1), // 27 spacer
1046 + Constraint::Length(1), // 28 Adapters heading
1047 + Constraint::Min(0), // 29 Adapters list
1048 ])
1049 .split(area);
1050
@@ -1085,10 +1087,73 @@ fn render_model_detail(frame: &mut Frame, app: &App, area: Rect) {
1087 rows[26],
1088 );
1089
1090 + // Adapters section
1091 + let adapter_heading = if app.adapter_list.is_empty() {
1092 + "LoRA Adapters"
1093 + } else {
1094 + "LoRA Adapters (↑↓ select · m merge & export)"
1095 + };
1096 frame.render_widget(
1089 - Paragraph::new("Esc · back to models").style(Style::new().fg(C_MUTED)),
1090 - rows[27],
1097 + Paragraph::new(adapter_heading).style(Style::new().fg(C_MUTED)),
1098 + rows[28],
1099 );
1100 + render_model_detail_adapters(frame, app, rows[29]);
1101 +}
1102 +
1103 +/// Render the selectable adapter list inside the Model Detail screen.
1104 +fn render_model_detail_adapters(frame: &mut Frame, app: &App, area: Rect) {
1105 + if app.adapter_list.is_empty() {
1106 + frame.render_widget(
1107 + Paragraph::new(" No adapters found. Press f to fine-tune.")
1108 + .style(Style::new().fg(C_MUTED)),
1109 + area,
1110 + );
1111 + return;
1112 + }
1113 +
1114 + let max_rows = area.height as usize;
1115 + for (i, adapter) in app.adapter_list.iter().enumerate().take(max_rows) {
1116 + let row_y = area.y + i as u16;
1117 + let row_area = Rect::new(area.x, row_y, area.width, 1);
1118 +
1119 + let is_selected = i == app.adapter_cursor;
1120 + let marker = if is_selected { "▸ " } else { " " };
1121 + let marker_style = if is_selected {
1122 + Style::new().fg(C_NEON).bold()
1123 + } else {
1124 + Style::new().fg(C_NEON)
1125 + };
1126 + let name_style = if is_selected {
1127 + Style::new().fg(C_TEXT).bold()
1128 + } else {
1129 + Style::new().fg(C_TEXT)
1130 + };
1131 + let meta_style = Style::new().fg(C_MUTED);
1132 +
1133 + // Truncate dir_name for display
1134 + let dir_display = if adapter.dir_name.len() > 16 {
1135 + format!("{}…", &adapter.dir_name[..15])
1136 + } else {
1137 + adapter.dir_name.clone()
1138 + };
1139 +
1140 + let line = Line::from(vec![
1141 + Span::styled(marker, marker_style),
1142 + Span::styled("◆ ", Style::new().fg(C_NEON)),
1143 + Span::styled(format!("{:<18}", dir_display), name_style),
1144 + Span::styled(format!("{:<10}", adapter.size), meta_style),
1145 + Span::styled(&adapter.modified, meta_style),
1146 + ]);
1147 +
1148 + if is_selected {
1149 + frame.render_widget(
1150 + Paragraph::new(line).style(Style::new().bg(C_SURFACE_STRONG)),
1151 + row_area,
1152 + );
1153 + } else {
1154 + frame.render_widget(Paragraph::new(line), row_area);
1155 + }
1156 + }
1157 }
1158
1159 fn fmt_bytes(bytes: i64) -> String {
@@ -1102,7 +1167,11 @@ fn fmt_bytes(bytes: i64) -> String {
1167 }
1168
1169 fn render_finetune(frame: &mut Frame, app: &App, area: Rect) {
1105 - if app.finetune_running || app.finetune_progress.is_some() {
1170 + if app.finetune_running
1171 + || app.merge_running
1172 + || app.gguf_running
1173 + || app.finetune_progress.is_some()
1174 + {
1175 render_finetune_progress(frame, app, area);
1176 } else {
1177 render_finetune_form(frame, app, area);
@@ -1367,7 +1436,7 @@ fn render_finetune_progress(frame: &mut Frame, app: &App, area: Rect) {
1436 }
1437
1438 // Hint line
1370 - let hint = if app.finetune_running {
1439 + let hint = if app.finetune_running || app.merge_running || app.gguf_running {
1440 "Ctrl+C · quit"
1441 } else {
1442 "Esc · back"
@@ -1375,7 +1444,8 @@ fn render_finetune_progress(frame: &mut Frame, app: &App, area: Rect) {
1444 frame.render_widget(Paragraph::new(hint).style(Style::new().fg(C_MUTED)), top[5]);
1445 }
1446
1378 -/// Render the fine-tune "Done" state with full adapter details.
1447 +/// Render the fine-tune "Done" state with full adapter details plus merge/export
1448 +/// progress and action hints.
1449 fn render_finetune_done(frame: &mut Frame, app: &App, adapter_path: &std::path::Path, area: Rect) {
1450 let rows = Layout::vertical([
1451 Constraint::Length(1), // 0 ✓ label
@@ -1388,8 +1458,7 @@ fn render_finetune_done(frame: &mut Frame, app: &App, adapter_path: &std::path::
1458 Constraint::Length(1), // 7 config heading
1459 Constraint::Length(1), // 8 rank + epochs + lr
1460 Constraint::Length(1), // 9 spacer
1391 - Constraint::Length(1), // 10 adapters heading
1392 - Constraint::Min(0), // 11 adapters list
1461 + Constraint::Min(0), // 10 merge/export section or adapters list
1462 ])
1463 .split(area);
1464
@@ -1447,12 +1516,159 @@ fn render_finetune_done(frame: &mut Frame, app: &App, adapter_path: &std::path::
1516 rows[8],
1517 );
1518
1450 - // Scan for all adapters under this model's snapshots directory
1519 + // Merge / GGUF export section
1520 + render_merge_gguf_section(frame, app, rows[10]);
1521 +}
1522 +
1523 +/// Render merge progress, GGUF export progress, and available action hints.
1524 +fn render_merge_gguf_section(frame: &mut Frame, app: &App, area: Rect) {
1525 + let rows = Layout::vertical([
1526 + Constraint::Length(1), // 0 merge status
1527 + Constraint::Length(1), // 1 spacer
1528 + Constraint::Length(1), // 2 gguf status
1529 + Constraint::Length(1), // 3 spacer
1530 + Constraint::Length(1), // 4 action hints
1531 + Constraint::Length(1), // 5 spacer
1532 + Constraint::Length(1), // 6 adapters heading
1533 + Constraint::Min(0), // 7 adapters list
1534 + ])
1535 + .split(area);
1536 +
1537 + // --- Merge status ---
1538 + match &app.merge_progress {
1539 + Some(crate::merge::MergeProgress::Loading) => {
1540 + frame.render_widget(
1541 + Paragraph::new("⠿ Loading model for merge…").style(Style::new().fg(C_MUTED)),
1542 + rows[0],
1543 + );
1544 + }
1545 + Some(crate::merge::MergeProgress::Merging { layer, total }) => {
1546 + frame.render_widget(
1547 + Paragraph::new(format!("⠿ Merging layer {layer}/{total}…"))
1548 + .style(Style::new().fg(C_MUTED)),
1549 + rows[0],
1550 + );
1551 + }
1552 + Some(crate::merge::MergeProgress::Saving) => {
1553 + frame.render_widget(
1554 + Paragraph::new("⠿ Saving merged model…").style(Style::new().fg(C_MUTED)),
1555 + rows[0],
1556 + );
1557 + }
1558 + Some(crate::merge::MergeProgress::Done { output_path }) => {
1559 + frame.render_widget(
1560 + Paragraph::new(Line::from(vec![
1561 + Span::styled("✓ ", Style::new().fg(C_NEON)),
1562 + Span::styled("Merged → ", Style::new().fg(C_NEON).bold()),
1563 + Span::styled(
1564 + output_path.to_string_lossy().to_string(),
1565 + Style::new().fg(C_TEXT),
1566 + ),
1567 + ])),
1568 + rows[0],
1569 + );
1570 + }
1571 + Some(crate::merge::MergeProgress::Failed(msg)) => {
1572 + frame.render_widget(
1573 + Paragraph::new(Line::from(vec![
1574 + Span::styled("✗ Merge failed: ", Style::new().fg(C_DANGER).bold()),
1575 + Span::styled(msg.as_str(), Style::new().fg(C_DANGER)),
1576 + ])),
1577 + rows[0],
1578 + );
1579 + }
1580 + None => {
1581 + frame.render_widget(
1582 + Paragraph::new(" Merge: not started").style(Style::new().fg(C_MUTED)),
1583 + rows[0],
1584 + );
1585 + }
1586 + }
1587 +
1588 + // --- GGUF export status ---
1589 + match &app.gguf_progress {
1590 + Some(crate::gguf::GgufProgress::ReadingModel) => {
1591 + frame.render_widget(
1592 + Paragraph::new("⠿ Reading model for GGUF export…").style(Style::new().fg(C_MUTED)),
1593 + rows[2],
1594 + );
1595 + }
1596 + Some(crate::gguf::GgufProgress::WritingTensor { index, total, name }) => {
1597 + frame.render_widget(
1598 + Paragraph::new(format!("⠿ Writing tensor {index}/{total} {name}"))
1599 + .style(Style::new().fg(C_MUTED)),
1600 + rows[2],
1601 + );
1602 + }
1603 + Some(crate::gguf::GgufProgress::Done {
1604 + output_path,
1605 + size_bytes,
1606 + }) => {
1607 + frame.render_widget(
1608 + Paragraph::new(Line::from(vec![
1609 + Span::styled("✓ ", Style::new().fg(C_NEON)),
1610 + Span::styled("GGUF → ", Style::new().fg(C_NEON).bold()),
1611 + Span::styled(
1612 + output_path.to_string_lossy().to_string(),
1613 + Style::new().fg(C_TEXT),
1614 + ),
1615 + Span::styled(
1616 + format!(" ({})", fmt_bytes(*size_bytes as i64)),
1617 + Style::new().fg(C_MUTED),
1618 + ),
1619 + ])),
1620 + rows[2],
1621 + );
1622 + }
1623 + Some(crate::gguf::GgufProgress::Failed(msg)) => {
1624 + frame.render_widget(
1625 + Paragraph::new(Line::from(vec![
1626 + Span::styled("✗ GGUF failed: ", Style::new().fg(C_DANGER).bold()),
1627 + Span::styled(msg.as_str(), Style::new().fg(C_DANGER)),
1628 + ])),
1629 + rows[2],
1630 + );
1631 + }
1632 + None => {
1633 + if app.merged_model_dir.is_some() {
1634 + frame.render_widget(
1635 + Paragraph::new(" GGUF: ready to export").style(Style::new().fg(C_MUTED)),
1636 + rows[2],
1637 + );
1638 + } else {
1639 + frame.render_widget(
1640 + Paragraph::new(" GGUF: merge first").style(Style::new().fg(C_MUTED)),
1641 + rows[2],
1642 + );
1643 + }
1644 + }
1645 + }
1646 +
1647 + // --- Action hints ---
1648 + let mut hints: Vec<Span> = Vec::new();
1649 + if !app.merge_running && !app.gguf_running {
1650 + hints.push(Span::styled("m", Style::new().fg(C_NEON)));
1651 + hints.push(Span::styled(
1652 + " · merge adapter ",
1653 + Style::new().fg(C_MUTED),
1654 + ));
1655 + if app.merged_model_dir.is_some() {
1656 + hints.push(Span::styled("g", Style::new().fg(C_NEON)));
1657 + hints.push(Span::styled(" · export GGUF ", Style::new().fg(C_MUTED)));
1658 + }
1659 + hints.push(Span::styled("Esc", Style::new().fg(C_NEON)));
1660 + hints.push(Span::styled(" · back", Style::new().fg(C_MUTED)));
1661 + }
1662 + if !hints.is_empty() {
1663 + frame.render_widget(Paragraph::new(Line::from(hints)), rows[4]);
1664 + }
1665 +
1666 + // --- Adapters list ---
1667 frame.render_widget(
1668 Paragraph::new("All Adapters").style(Style::new().fg(C_MUTED)),
1453 - rows[10],
1669 + rows[6],
1670 );
1455 - render_adapter_list(frame, app, rows[11]);
1671 + render_adapter_list(frame, app, rows[7]);
1672 }
1673
1674 /// Scan the model's cache directory for existing LoRA adapter files and list them.
@@ -1615,10 +1831,30 @@ fn render_footer(frame: &mut Frame, app: &App, area: Rect) {
1831 Span::styled("Ctrl+C", Style::new().fg(C_NEON)),
1832 Span::styled(" · quit", Style::new().fg(C_MUTED)),
1833 ],
1618 - Screen::FineTune if app.finetune_running => vec![
1834 + Screen::FineTune if app.finetune_running || app.merge_running || app.gguf_running => vec![
1835 Span::styled("Ctrl+C", Style::new().fg(C_NEON)),
1836 Span::styled(" · quit", Style::new().fg(C_MUTED)),
1837 ],
1838 + Screen::FineTune
1839 + if matches!(
1840 + &app.finetune_progress,
1841 + Some(crate::finetune::FineTuneProgress::Done { .. })
1842 + ) =>
1843 + {
1844 + let mut keys = vec![
1845 + Span::styled("m", Style::new().fg(C_NEON)),
1846 + Span::styled(" · merge ", Style::new().fg(C_MUTED)),
1847 + ];
1848 + if app.merged_model_dir.is_some() {
1849 + keys.push(Span::styled("g", Style::new().fg(C_NEON)));
1850 + keys.push(Span::styled(" · export GGUF ", Style::new().fg(C_MUTED)));
1851 + }
1852 + keys.push(Span::styled("Esc", Style::new().fg(C_NEON)));
1853 + keys.push(Span::styled(" · back ", Style::new().fg(C_MUTED)));
1854 + keys.push(Span::styled("Ctrl+C", Style::new().fg(C_NEON)));
1855 + keys.push(Span::styled(" · quit", Style::new().fg(C_MUTED)));
1856 + keys
1857 + }
1858 Screen::FineTune => vec![
1859 Span::styled("Tab", Style::new().fg(C_NEON)),
1860 Span::styled(" · next field ", Style::new().fg(C_MUTED)),
@@ -1629,14 +1865,23 @@ fn render_footer(frame: &mut Frame, app: &App, area: Rect) {
1865 Span::styled("Ctrl+C", Style::new().fg(C_NEON)),
1866 Span::styled(" · quit", Style::new().fg(C_MUTED)),
1867 ],
1632 - Screen::ModelDetail => vec![
1633 - Span::styled("f", Style::new().fg(C_NEON)),
1634 - Span::styled(" · fine-tune ", Style::new().fg(C_MUTED)),
1635 - Span::styled("Esc", Style::new().fg(C_NEON)),
1636 - Span::styled(" · back ", Style::new().fg(C_MUTED)),
1637 - Span::styled("Ctrl+C", Style::new().fg(C_NEON)),
1638 - Span::styled(" · quit", Style::new().fg(C_MUTED)),
1639 - ],
1868 + Screen::ModelDetail => {
1869 + let mut keys = Vec::new();
1870 + if !app.adapter_list.is_empty() {
1871 + keys.push(Span::styled("Enter", Style::new().fg(C_NEON)));
1872 + keys.push(Span::styled(
1873 + " · merge & export ",
1874 + Style::new().fg(C_MUTED),
1875 + ));
1876 + }
1877 + keys.push(Span::styled("f", Style::new().fg(C_NEON)));
1878 + keys.push(Span::styled(" · fine-tune ", Style::new().fg(C_MUTED)));
1879 + keys.push(Span::styled("Esc", Style::new().fg(C_NEON)));
1880 + keys.push(Span::styled(" · back ", Style::new().fg(C_MUTED)));
1881 + keys.push(Span::styled("Ctrl+C", Style::new().fg(C_NEON)));
1882 + keys.push(Span::styled(" · quit", Style::new().fg(C_MUTED)));
1883 + keys
1884 + }
1885 };
1886
1887 frame.render_widget(