Skip to main content

simple_chat/
simple_chat.rs

1use litert_lm::{Backend, Engine};
2use std::io::{self, Write};
3
4fn main() -> Result<(), Box<dyn std::error::Error>> {
5    // Get model path from command line argument
6    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    // Create engine with CPU backend
17    let engine = Engine::new(model_path, Backend::Cpu)?;
18    println!("Engine created successfully!");
19
20    // Create a session (conversation)
21    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    // Interactive chat loop
29    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        // Generate response
47        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}