1use crate::state::{HistoryEntry, ModelRole, SessionState};
4use entrenar_common::{EntrenarError, Result};
5
6#[derive(Debug, Clone, PartialEq)]
8pub enum Command {
9 Fetch { model_id: String, role: ModelRole },
11 Inspect { target: InspectTarget },
13 Memory {
15 batch_size: Option<u32>,
16 seq_len: Option<usize>,
17 },
18 Set { key: String, value: String },
20 Distill { dry_run: bool },
22 Export { format: String, path: String },
24 History,
26 Help { topic: Option<String> },
28 Clear,
30 Quit,
32 Unknown { input: String },
34}
35
36#[derive(Debug, Clone, PartialEq)]
38pub enum InspectTarget {
39 Layers,
41 Memory,
43 All,
45 Model(String),
47}
48
49pub fn parse(input: &str) -> Result<Command> {
51 let input = input.trim();
52 if input.is_empty() {
53 return Ok(Command::Unknown {
54 input: String::new(),
55 });
56 }
57
58 let parts: Vec<&str> = input.split_whitespace().collect();
59 let cmd = parts[0].to_lowercase();
60 let args = &parts[1..];
61
62 match cmd.as_str() {
63 "fetch" | "download" => parse_fetch(args),
64 "inspect" | "show" => parse_inspect(args),
65 "memory" | "mem" => parse_memory(args),
66 "set" => parse_set(args),
67 "distill" | "train" => parse_distill(args),
68 "export" | "save" => parse_export(args),
69 "history" | "hist" => Ok(Command::History),
70 "help" | "?" => parse_help(args),
71 "clear" | "cls" => Ok(Command::Clear),
72 "quit" | "exit" | "q" => Ok(Command::Quit),
73 _ => Ok(Command::Unknown {
74 input: input.to_string(),
75 }),
76 }
77}
78
79fn parse_fetch(args: &[&str]) -> Result<Command> {
80 if args.is_empty() {
81 return Err(EntrenarError::ConfigValue {
82 field: "model_id".into(),
83 message: "No model ID provided".into(),
84 suggestion: "Usage: fetch <model_id> [--teacher|--student]".into(),
85 });
86 }
87
88 let model_id = args[0].to_string();
89 let role = if args.contains(&"--teacher") {
90 ModelRole::Teacher
91 } else if args.contains(&"--student") {
92 ModelRole::Student
93 } else {
94 ModelRole::None
95 };
96
97 Ok(Command::Fetch { model_id, role })
98}
99
100fn parse_inspect(args: &[&str]) -> Result<Command> {
101 let target = if args.is_empty() {
102 InspectTarget::All
103 } else {
104 match args[0].to_lowercase().as_str() {
105 "layers" | "layer" => InspectTarget::Layers,
106 "memory" | "mem" => InspectTarget::Memory,
107 "all" => InspectTarget::All,
108 name => InspectTarget::Model(name.to_string()),
109 }
110 };
111
112 Ok(Command::Inspect { target })
113}
114
115fn parse_memory(args: &[&str]) -> Result<Command> {
116 let mut batch_size = None;
117 let mut seq_len = None;
118
119 let mut i = 0;
120 while i < args.len() {
121 match args[i] {
122 "--batch" | "-b" if i + 1 < args.len() => {
123 batch_size = args[i + 1].parse().ok();
124 i += 2;
125 }
126 "--seq" | "-s" if i + 1 < args.len() => {
127 seq_len = args[i + 1].parse().ok();
128 i += 2;
129 }
130 _ => i += 1,
131 }
132 }
133
134 Ok(Command::Memory {
135 batch_size,
136 seq_len,
137 })
138}
139
140fn parse_set(args: &[&str]) -> Result<Command> {
141 if args.len() < 2 {
142 return Err(EntrenarError::ConfigValue {
143 field: "set".into(),
144 message: "Not enough arguments".into(),
145 suggestion: "Usage: set <key> <value>".into(),
146 });
147 }
148
149 Ok(Command::Set {
150 key: args[0].to_string(),
151 value: args[1..].join(" "),
152 })
153}
154
155fn parse_distill(args: &[&str]) -> Result<Command> {
156 let dry_run = args.contains(&"--dry-run") || args.contains(&"-n");
157 Ok(Command::Distill { dry_run })
158}
159
160fn parse_export(args: &[&str]) -> Result<Command> {
161 if args.len() < 2 {
162 return Err(EntrenarError::ConfigValue {
163 field: "export".into(),
164 message: "Not enough arguments".into(),
165 suggestion: "Usage: export <format> <path>".into(),
166 });
167 }
168
169 Ok(Command::Export {
170 format: args[0].to_string(),
171 path: args[1].to_string(),
172 })
173}
174
175fn parse_help(args: &[&str]) -> Result<Command> {
176 let topic = args.first().map(ToString::to_string);
177 Ok(Command::Help { topic })
178}
179
180pub fn execute(cmd: &Command, state: &mut SessionState) -> Result<String> {
182 let start = std::time::Instant::now();
183
184 let result = match cmd {
185 Command::Fetch { model_id, .. } => execute_fetch(model_id),
189 Command::Inspect { target } => execute_inspect(target, state),
190 Command::Memory {
191 batch_size,
192 seq_len,
193 } => execute_memory(*batch_size, *seq_len, state),
194 Command::Set { key, value } => execute_set(key, value, state),
195 Command::Distill { dry_run } => execute_distill(*dry_run, state),
196 Command::Export { format, path } => execute_export(format, path, state),
197 Command::History => execute_history(state),
198 Command::Help { topic } => execute_help(topic.as_deref()),
199 Command::Clear => Ok(String::new()),
200 Command::Quit => Ok("Goodbye!".to_string()),
201 Command::Unknown { input } => {
202 if input.is_empty() {
203 Ok(String::new())
204 } else {
205 Err(EntrenarError::ConfigValue {
206 field: "command".into(),
207 message: format!("Unknown command: {input}"),
208 suggestion: "Type 'help' for available commands".into(),
209 })
210 }
211 }
212 };
213
214 let duration_ms = start.elapsed().as_millis() as u64;
215 let success = result.is_ok();
216
217 if !matches!(
219 cmd,
220 Command::Help { .. }
221 | Command::History
222 | Command::Clear
223 | Command::Quit
224 | Command::Unknown { .. }
225 ) {
226 let cmd_str = format!("{cmd:?}");
227 state.add_to_history(HistoryEntry::new(cmd_str, duration_ms, success));
228 state.record_command(duration_ms, success);
229 }
230
231 result
232}
233
234fn execute_fetch(model_id: &str) -> Result<String> {
235 Err(EntrenarError::ConfigValue {
263 field: "fetch".into(),
264 message: format!(
265 "cannot fetch `{model_id}`: this shell has no HuggingFace client, so it \
266 downloads nothing. It previously reported success for any string at all, \
267 with a parameter count and layer count string-matched out of the model \
268 ID itself"
269 ),
270 suggestion: "Download with `apr pull <model_id>` or `apr import hf://<model_id>`, \
271 then read the real file with `apr inspect` / `apr tensors`. \
272 Tracked in #2519."
273 .into(),
274 })
275}
276
277fn execute_inspect(target: &InspectTarget, state: &SessionState) -> Result<String> {
278 match target {
279 InspectTarget::All => {
280 if state.loaded_models().is_empty() {
281 return Ok("No models loaded. Use 'fetch <model_id>' to load a model.".to_string());
282 }
283
284 let mut output = String::from("Loaded Models:\n");
285 for (name, model) in state.loaded_models() {
286 output.push_str(&format!(
287 " {} ({}): {:.1}B params, {} layers\n",
288 name,
289 model.id,
290 model.parameters as f64 / 1e9,
291 model.layers
292 ));
293 }
294 Ok(output)
295 }
296 InspectTarget::Layers => {
297 let mut output = String::from("Layer Analysis:\n");
298 for (name, model) in state.loaded_models() {
299 output.push_str(&format!(
300 " {}: {} layers, hidden_dim={}\n",
301 name, model.layers, model.hidden_dim
302 ));
303 }
304 Ok(output)
305 }
306 InspectTarget::Memory => execute_memory(None, None, state),
307 InspectTarget::Model(name) => {
308 if let Some(model) = state.get_model(name) {
309 Ok(format!(
310 "Model: {}\n ID: {}\n Path: {}\n Architecture: {}\n Parameters: {:.1}B\n Layers: {}\n Hidden Dim: {}",
311 name, model.id, model.path.display(), model.architecture,
312 model.parameters as f64 / 1e9, model.layers, model.hidden_dim
313 ))
314 } else {
315 Err(EntrenarError::ModelNotFound { path: name.into() })
316 }
317 }
318 }
319}
320
321fn execute_memory(
322 batch_size: Option<u32>,
323 seq_len: Option<usize>,
324 state: &SessionState,
325) -> Result<String> {
326 let batch = batch_size.unwrap_or(state.preferences().default_batch_size);
327 let seq = seq_len.unwrap_or(state.preferences().default_seq_len);
328
329 let total_params: u64 = state.loaded_models().values().map(|m| m.parameters).sum();
330 let model_mem = total_params * 2; let activation_mem = u64::from(batch) * (seq as u64) * 4096 * 32 * 2;
332 let total = model_mem + activation_mem;
333
334 Ok(format!(
335 "Memory Estimate (batch={}, seq={}):\n Model: {:.1} GB\n Activations: {:.1} GB\n Total: {:.1} GB",
336 batch, seq,
337 model_mem as f64 / 1e9,
338 activation_mem as f64 / 1e9,
339 total as f64 / 1e9
340 ))
341}
342
343fn execute_set(key: &str, value: &str, state: &mut SessionState) -> Result<String> {
344 match key {
345 "batch_size" | "batch" => {
346 let v: u32 = value.parse().map_err(|_| EntrenarError::ConfigValue {
347 field: "batch_size".into(),
348 message: "Invalid number".into(),
349 suggestion: "Use a positive integer".into(),
350 })?;
351 state.preferences_mut().default_batch_size = v;
352 Ok(format!("Set batch_size = {v}"))
353 }
354 "seq_len" | "seq" => {
355 let v: usize = value.parse().map_err(|_| EntrenarError::ConfigValue {
356 field: "seq_len".into(),
357 message: "Invalid number".into(),
358 suggestion: "Use a positive integer".into(),
359 })?;
360 state.preferences_mut().default_seq_len = v;
361 Ok(format!("Set seq_len = {v}"))
362 }
363 _ => Err(EntrenarError::ConfigValue {
364 field: key.into(),
365 message: "Unknown setting".into(),
366 suggestion: "Available settings: batch_size, seq_len".into(),
367 }),
368 }
369}
370
371fn execute_distill(dry_run: bool, state: &SessionState) -> Result<String> {
372 let teacher = state
373 .loaded_models()
374 .values()
375 .find(|m| m.role == ModelRole::Teacher);
376 let student = state
377 .loaded_models()
378 .values()
379 .find(|m| m.role == ModelRole::Student);
380
381 if teacher.is_none() {
382 return Err(EntrenarError::ConfigValue {
383 field: "teacher".into(),
384 message: "No teacher model loaded".into(),
385 suggestion: "Use 'fetch <model_id> --teacher' to load a teacher model".into(),
386 });
387 }
388
389 if student.is_none() {
390 return Err(EntrenarError::ConfigValue {
391 field: "student".into(),
392 message: "No student model loaded".into(),
393 suggestion: "Use 'fetch <model_id> --student' to load a student model".into(),
394 });
395 }
396
397 if dry_run {
398 let teacher = teacher.expect("teacher validated non-None above");
399 let student = student.expect("student validated non-None above");
400 Ok(format!(
401 "Dry run configuration:\n Teacher: {} ({:.1}B)\n Student: {} ({:.1}B)\n Ready to train",
402 teacher.id, teacher.parameters as f64 / 1e9,
403 student.id, student.parameters as f64 / 1e9
404 ))
405 } else {
406 Err(EntrenarError::ConfigValue {
415 field: "distill".into(),
416 message: format!(
417 "cannot distill `{}` into `{}`: this shell has no training loop, \
422 no optimizer and no dataset. It previously reported that \
423 training had begun, and exited 0 without training anything",
424 teacher.map_or("<teacher>", |t| t.id.as_str()),
425 student.map_or("<student>", |s| s.id.as_str()),
426 ),
427 suggestion: "Run the real trainer: `apr distill` / `apr finetune`. \
428 `distill --dry-run` here still prints the configuration. \
429 Tracked in #2519."
430 .into(),
431 })
432 }
433}
434
435fn execute_export(format: &str, path: &str, _state: &SessionState) -> Result<String> {
436 Err(EntrenarError::ConfigValue {
444 field: "export".into(),
445 message: format!("export is not implemented in this shell ({format} -> {path})"),
446 suggestion: "Use `apr export` for real format conversion.".into(),
447 })
448}
449
450fn execute_history(state: &SessionState) -> Result<String> {
451 if state.history().is_empty() {
452 return Ok("No command history.".to_string());
453 }
454
455 let mut output = String::from("Command History:\n");
456 for (i, entry) in state.history().iter().enumerate() {
457 let status = if entry.success { "✓" } else { "✗" };
458 output.push_str(&format!(
459 " {}. {} {} ({}ms)\n",
460 i + 1,
461 status,
462 entry.command,
463 entry.duration_ms
464 ));
465 }
466 Ok(output)
467}
468
469fn execute_help(topic: Option<&str>) -> Result<String> {
470 match topic {
471 Some("fetch") => Ok(
472 "fetch <model_id> [--teacher|--student]\n Download a model from HuggingFace"
473 .to_string(),
474 ),
475 Some("inspect") => {
476 Ok("inspect [layers|memory|all|<model>]\n Inspect loaded models".to_string())
477 }
478 Some("memory") => {
479 Ok("memory [--batch <n>] [--seq <n>]\n Estimate memory requirements".to_string())
480 }
481 Some("distill") => Ok("distill [--dry-run]\n Run distillation training".to_string()),
482 _ => Ok("Available commands:
483 fetch <model> Download model from HuggingFace
484 inspect [target] Inspect loaded models
485 memory Estimate memory requirements
486 set <key> <value> Configure settings
487 distill Run distillation
488 export <fmt> <path> Export model
489 history Show command history
490 help [topic] Show help
491 quit Exit shell"
492 .to_string()),
493 }
494}
495
496#[cfg(test)]
507const ARCH_PATTERNS: &[(&[&str], &str)] = &[
508 (&["qwen"], "qwen"),
509 (&["phi"], "phi"),
510 (&["falcon"], "falcon"),
511 (&["mistral", "mixtral"], "mistral"),
512 (&["llama"], "llama"),
513 (&["bert"], "bert"),
514 (&["gpt"], "gpt"),
515];
516
517#[cfg(test)]
518fn detect_architecture(model_id: &str) -> String {
519 let lower = model_id.to_lowercase();
520 for (patterns, arch) in ARCH_PATTERNS {
521 if patterns.iter().any(|p| lower.contains(p)) {
522 return (*arch).to_string();
523 }
524 }
525 eprintln!(
526 "Warning: could not detect architecture from model ID '{model_id}', \
527 defaulting to 'unknown' (use tensor-based detection for accuracy)"
528 );
529 "unknown".to_string()
530}
531
532#[cfg(test)]
533fn estimate_params(model_id: &str) -> u64 {
534 let lower = model_id.to_lowercase();
535 if lower.contains("70b") {
536 70_000_000_000
537 } else if lower.contains("13b") {
538 13_000_000_000
539 } else if lower.contains("7b") {
540 7_000_000_000
541 } else if lower.contains("1.1b") || lower.contains("1b") {
542 1_100_000_000
543 } else if lower.contains("base") {
544 350_000_000
545 } else {
546 1_000_000_000
547 }
548}
549
550#[cfg(test)]
551fn estimate_layers(model_id: &str) -> u32 {
552 let lower = model_id.to_lowercase();
553 if lower.contains("70b") {
554 80
555 } else if lower.contains("13b") {
556 40
557 } else if lower.contains("7b") {
558 32
559 } else if lower.contains("base") {
560 12
561 } else {
562 24
563 }
564}
565
566#[cfg(test)]
567mod tests {
568 use super::*;
569 use crate::state::LoadedModel;
571
572 #[test]
573 fn test_parse_fetch() {
574 let cmd = parse("fetch meta-llama/Llama-2-7b --teacher").expect("parsing should succeed");
575 assert!(matches!(
576 cmd,
577 Command::Fetch {
578 role: ModelRole::Teacher,
579 ..
580 }
581 ));
582
583 let cmd =
584 parse("fetch TinyLlama/TinyLlama-1.1B --student").expect("parsing should succeed");
585 assert!(matches!(
586 cmd,
587 Command::Fetch {
588 role: ModelRole::Student,
589 ..
590 }
591 ));
592 }
593
594 #[test]
595 fn test_parse_inspect() {
596 assert!(matches!(
597 parse("inspect").expect("parsing should succeed"),
598 Command::Inspect {
599 target: InspectTarget::All
600 }
601 ));
602 assert!(matches!(
603 parse("inspect layers").expect("parsing should succeed"),
604 Command::Inspect {
605 target: InspectTarget::Layers
606 }
607 ));
608 }
609
610 #[test]
611 fn test_parse_memory() {
612 let cmd = parse("memory --batch 64 --seq 1024").expect("parsing should succeed");
613 if let Command::Memory {
614 batch_size,
615 seq_len,
616 } = cmd
617 {
618 assert_eq!(batch_size, Some(64));
619 assert_eq!(seq_len, Some(1024));
620 } else {
621 panic!("Expected Memory command");
622 }
623 }
624
625 #[test]
626 fn test_parse_quit_variants() {
627 assert!(matches!(
628 parse("quit").expect("parsing should succeed"),
629 Command::Quit
630 ));
631 assert!(matches!(
632 parse("exit").expect("parsing should succeed"),
633 Command::Quit
634 ));
635 assert!(matches!(
636 parse("q").expect("parsing should succeed"),
637 Command::Quit
638 ));
639 }
640
641 #[test]
645 fn test_execute_fetch_refuses_and_loads_nothing() {
646 let mut state = SessionState::new();
647 let err = execute(
648 &Command::Fetch {
649 model_id: "meta-llama/Llama-2-7b".to_string(),
650 role: ModelRole::Teacher,
651 },
652 &mut state,
653 )
654 .expect_err("fetch must not claim to have downloaded a model");
655
656 assert!(format!("{err}").contains("no HuggingFace client"));
657 assert!(state.get_model("teacher").is_none());
658 assert!(state.loaded_models().is_empty());
659 }
660
661 #[test]
662 fn test_execute_set() {
663 let mut state = SessionState::new();
664
665 execute_set("batch_size", "64", &mut state).expect("operation should succeed");
666 assert_eq!(state.preferences().default_batch_size, 64);
667
668 execute_set("seq_len", "1024", &mut state).expect("operation should succeed");
669 assert_eq!(state.preferences().default_seq_len, 1024);
670 }
671
672 #[test]
673 fn test_unknown_command() {
674 let cmd = parse("foobar").expect("parsing should succeed");
675 assert!(matches!(cmd, Command::Unknown { .. }));
676 }
677
678 #[test]
679 fn test_parse_fetch_missing_model() {
680 let result = parse("fetch");
681 assert!(result.is_err());
682 }
683
684 #[test]
685 fn test_parse_set_not_enough_args() {
686 let result = parse("set batch_size");
687 assert!(result.is_err());
688 }
689
690 #[test]
691 fn test_parse_export_not_enough_args() {
692 let result = parse("export safetensors");
693 assert!(result.is_err());
694 }
695
696 #[test]
697 fn test_parse_export_valid() {
698 let cmd = parse("export safetensors /tmp/model.st").expect("parsing should succeed");
699 if let Command::Export { format, path } = cmd {
700 assert_eq!(format, "safetensors");
701 assert_eq!(path, "/tmp/model.st");
702 } else {
703 panic!("Expected Export command");
704 }
705 }
706
707 #[test]
708 fn test_parse_help_with_topic() {
709 let cmd = parse("help fetch").expect("parsing should succeed");
710 if let Command::Help { topic } = cmd {
711 assert_eq!(topic, Some("fetch".to_string()));
712 } else {
713 panic!("Expected Help command");
714 }
715 }
716
717 #[test]
718 fn test_parse_distill_dry_run() {
719 let cmd = parse("distill --dry-run").expect("parsing should succeed");
720 if let Command::Distill { dry_run } = cmd {
721 assert!(dry_run);
722 } else {
723 panic!("Expected Distill command");
724 }
725 }
726
727 #[test]
728 fn test_parse_distill_short_flag() {
729 let cmd = parse("distill -n").expect("parsing should succeed");
730 if let Command::Distill { dry_run } = cmd {
731 assert!(dry_run);
732 } else {
733 panic!("Expected Distill command");
734 }
735 }
736
737 #[test]
738 fn test_parse_inspect_model() {
739 let cmd = parse("inspect teacher").expect("parsing should succeed");
740 if let Command::Inspect { target } = cmd {
741 assert_eq!(target, InspectTarget::Model("teacher".to_string()));
742 } else {
743 panic!("Expected Inspect command");
744 }
745 }
746
747 #[test]
748 fn test_parse_inspect_memory() {
749 let cmd = parse("inspect memory").expect("parsing should succeed");
750 assert!(matches!(
751 cmd,
752 Command::Inspect {
753 target: InspectTarget::Memory
754 }
755 ));
756 }
757
758 #[test]
759 fn test_parse_command_aliases() {
760 assert!(matches!(
762 parse("download model").expect("parsing should succeed"),
763 Command::Fetch { .. }
764 ));
765 assert!(matches!(
767 parse("show layers").expect("parsing should succeed"),
768 Command::Inspect { .. }
769 ));
770 assert!(matches!(
772 parse("mem").expect("parsing should succeed"),
773 Command::Memory { .. }
774 ));
775 assert!(matches!(
777 parse("train").expect("parsing should succeed"),
778 Command::Distill { .. }
779 ));
780 assert!(matches!(
782 parse("save gguf /tmp/out").expect("parsing should succeed"),
783 Command::Export { .. }
784 ));
785 assert!(matches!(
787 parse("cls").expect("parsing should succeed"),
788 Command::Clear
789 ));
790 assert!(matches!(
792 parse("?").expect("parsing should succeed"),
793 Command::Help { .. }
794 ));
795 assert!(matches!(
797 parse("hist").expect("parsing should succeed"),
798 Command::History
799 ));
800 }
801
802 #[test]
803 fn test_execute_inspect_no_models() {
804 let state = SessionState::new();
805 let result = execute_inspect(&InspectTarget::All, &state);
806 assert!(result
807 .expect("load should succeed")
808 .contains("No models loaded"));
809 }
810
811 #[test]
812 fn test_execute_inspect_layers() {
813 let mut state = SessionState::new();
814 let model = LoadedModel {
815 id: "test".to_string(),
816 path: std::path::PathBuf::from("/tmp"),
817 architecture: "llama".to_string(),
818 parameters: 7_000_000_000,
819 layers: 32,
820 hidden_dim: 4096,
821 role: ModelRole::None,
822 };
823 state.add_model("test".to_string(), model);
824
825 let result =
826 execute_inspect(&InspectTarget::Layers, &state).expect("operation should succeed");
827 assert!(result.contains("32 layers"));
828 }
829
830 #[test]
831 fn test_execute_inspect_model_not_found() {
832 let state = SessionState::new();
833 let result = execute_inspect(&InspectTarget::Model("unknown".to_string()), &state);
834 assert!(result.is_err());
835 }
836
837 #[test]
838 fn test_execute_history_empty() {
839 let state = SessionState::new();
840 let result = execute_history(&state).expect("operation should succeed");
841 assert!(result.contains("No command history"));
842 }
843
844 #[test]
845 fn test_execute_help_topics() {
846 let fetch_help = execute_help(Some("fetch")).expect("operation should succeed");
847 assert!(fetch_help.contains("Download"));
848
849 let inspect_help = execute_help(Some("inspect")).expect("operation should succeed");
850 assert!(inspect_help.contains("Inspect"));
851
852 let memory_help = execute_help(Some("memory")).expect("operation should succeed");
853 assert!(memory_help.contains("memory"));
854
855 let distill_help = execute_help(Some("distill")).expect("operation should succeed");
856 assert!(distill_help.contains("distill"));
857
858 let general_help = execute_help(None).expect("operation should succeed");
859 assert!(general_help.contains("Available commands"));
860 }
861
862 #[test]
863 fn test_detect_architecture_variants() {
864 assert_eq!(detect_architecture("meta-llama/Llama-2-7b"), "llama");
865 assert_eq!(detect_architecture("bert-base-uncased"), "bert");
866 assert_eq!(detect_architecture("openai-gpt"), "gpt");
867 assert_eq!(detect_architecture("mistralai/Mistral-7B"), "mistral");
868 assert_eq!(detect_architecture("Qwen/Qwen2.5-Coder-0.5B"), "qwen");
869 assert_eq!(detect_architecture("microsoft/phi-2"), "phi");
870 assert_eq!(detect_architecture("tiiuae/falcon-7b"), "falcon");
871 assert_eq!(detect_architecture("mistralai/Mixtral-8x7B"), "mistral");
872 assert_eq!(detect_architecture("custom-model"), "unknown");
873 }
874
875 #[test]
880 fn test_falsify_n04_mistral_before_llama_in_ambiguous_id() {
881 assert_eq!(
885 detect_architecture("mistral-llama-variant"),
886 "mistral",
887 "Mistral should take priority over LLaMA in ambiguous IDs"
888 );
889 }
890
891 #[test]
892 fn test_falsify_n04_unknown_is_explicit() {
893 let result = detect_architecture("totally-custom-model-v3");
896 assert_eq!(result, "unknown");
897 assert!(!result.is_empty());
898 }
899
900 #[test]
901 fn test_estimate_params_variants() {
902 assert_eq!(estimate_params("model-70b"), 70_000_000_000);
903 assert_eq!(estimate_params("model-13b"), 13_000_000_000);
904 assert_eq!(estimate_params("model-7b"), 7_000_000_000);
905 assert_eq!(estimate_params("model-1.1b"), 1_100_000_000);
906 assert_eq!(estimate_params("bert-base"), 350_000_000);
907 }
908
909 #[test]
910 fn test_estimate_layers_variants() {
911 assert_eq!(estimate_layers("model-70b"), 80);
912 assert_eq!(estimate_layers("model-13b"), 40);
913 assert_eq!(estimate_layers("model-7b"), 32);
914 assert_eq!(estimate_layers("bert-base"), 12);
915 }
916
917 #[test]
918 fn test_execute_set_invalid_number() {
919 let mut state = SessionState::new();
920 let result = execute_set("batch_size", "not_a_number", &mut state);
921 assert!(result.is_err());
922 }
923
924 #[test]
925 fn test_execute_set_unknown_key() {
926 let mut state = SessionState::new();
927 let result = execute_set("unknown_setting", "value", &mut state);
928 assert!(result.is_err());
929 }
930
931 #[test]
932 fn test_execute_distill_no_teacher() {
933 let state = SessionState::new();
934 let result = execute_distill(true, &state);
935 assert!(result.is_err());
936 }
937
938 #[test]
939 fn test_execute_distill_no_student() {
940 let mut state = SessionState::new();
941 let model = LoadedModel {
942 id: "teacher".to_string(),
943 path: std::path::PathBuf::from("/tmp"),
944 architecture: "llama".to_string(),
945 parameters: 7_000_000_000,
946 layers: 32,
947 hidden_dim: 4096,
948 role: ModelRole::Teacher,
949 };
950 state.add_model("teacher".to_string(), model);
951
952 let result = execute_distill(true, &state);
953 assert!(result.is_err());
954 }
955
956 #[test]
957 fn test_execute_distill_success() {
958 let mut state = SessionState::new();
959
960 let teacher = LoadedModel {
961 id: "teacher/model".to_string(),
962 path: std::path::PathBuf::from("/tmp/t"),
963 architecture: "llama".to_string(),
964 parameters: 7_000_000_000,
965 layers: 32,
966 hidden_dim: 4096,
967 role: ModelRole::Teacher,
968 };
969 state.add_model("teacher".to_string(), teacher);
970
971 let student = LoadedModel {
972 id: "student/model".to_string(),
973 path: std::path::PathBuf::from("/tmp/s"),
974 architecture: "llama".to_string(),
975 parameters: 1_000_000_000,
976 layers: 12,
977 hidden_dim: 2048,
978 role: ModelRole::Student,
979 };
980 state.add_model("student".to_string(), student);
981
982 let result = execute_distill(true, &state).expect("operation should succeed");
983 assert!(result.contains("Dry run"));
984 }
985
986 #[test]
987 fn test_execute_export_refuses_rather_than_fabricating() {
988 let state = SessionState::new();
989 let result = execute_export("safetensors", "/tmp/model.st", &state);
990 assert!(
995 result.is_err(),
996 "export must not report success without exporting"
997 );
998 let msg = format!("{}", result.unwrap_err());
999 assert!(
1000 !msg.contains("Exported to"),
1001 "the error must not read like a success: {msg}"
1002 );
1003 assert!(msg.contains("not implemented"), "must say why: {msg}");
1004 }
1005
1006 #[test]
1007 fn test_execute_export_leaves_no_file_behind() {
1008 let state = SessionState::new();
1011 let path = std::env::temp_dir().join("apr-2519-export-probe.st");
1012 let _ = std::fs::remove_file(&path);
1013 let _ = execute_export("safetensors", path.to_str().unwrap(), &state);
1014 assert!(!path.exists(), "export refused but still created {path:?}");
1015 }
1016
1017 #[test]
1018 fn test_parse_empty_input() {
1019 let cmd = parse("").expect("parsing should succeed");
1020 assert!(matches!(cmd, Command::Unknown { .. }));
1021 }
1022
1023 #[test]
1024 fn test_execute_memory_with_args() {
1025 let state = SessionState::new();
1026 let result =
1027 execute_memory(Some(64), Some(1024), &state).expect("operation should succeed");
1028 assert!(result.contains("batch=64"));
1029 assert!(result.contains("seq=1024"));
1030 }
1031
1032 #[test]
1033 fn test_command_enum_equality() {
1034 assert_eq!(Command::Quit, Command::Quit);
1035 assert_eq!(Command::Clear, Command::Clear);
1036 assert_eq!(Command::History, Command::History);
1037 }
1038}