simple_chat/
simple_chat.rs1use litert_lm::{Backend, Engine};
2use std::io::{self, Write};
3
4fn main() -> Result<(), Box<dyn std::error::Error>> {
5 let args: Vec<String> = std::env::args().collect();
7 if args.len() < 2 {
8 eprintln!("Usage: {} <model_path>", args[0]);
9 eprintln!("Example: {} model.tflite", args[0]);
10 std::process::exit(1);
11 }
12 let model_path = &args[1];
13
14 println!("Loading model from: {}", model_path);
15
16 let engine = Engine::new(model_path, Backend::Cpu)?;
18 println!("Engine created successfully!");
19
20 let session = engine.create_session()?;
22 println!("Session created successfully!");
23 println!();
24 println!("You can now chat with the model. Type 'quit' or 'exit' to stop.");
25 println!("========================================");
26 println!();
27
28 loop {
30 print!("You: ");
31 io::stdout().flush()?;
32
33 let mut input = String::new();
34 io::stdin().read_line(&mut input)?;
35 let input = input.trim();
36
37 if input.is_empty() {
38 continue;
39 }
40
41 if input.eq_ignore_ascii_case("quit") || input.eq_ignore_ascii_case("exit") {
42 println!("Goodbye!");
43 break;
44 }
45
46 print!("Assistant: ");
48 io::stdout().flush()?;
49
50 match session.generate(input) {
51 Ok(response) => {
52 println!("{}", response);
53 println!();
54 }
55 Err(e) => {
56 eprintln!("Error generating response: {}", e);
57 println!();
58 }
59 }
60 }
61
62 Ok(())
63}