Skip to main content

imatrix/
main.rs

1use anyhow::Result;
2use mistralrs::{
3    IsqType, PagedAttentionMetaBuilder, TextMessageRole, TextMessages, TextModelBuilder,
4};
5
6#[tokio::main]
7async fn main() -> Result<()> {
8    let model = TextModelBuilder::new("meta-llama/Llama-3.2-3B-Instruct")
9        .with_isq(IsqType::Q4K)
10        .with_calibration_file("calibration_data/calibration_datav3_small.txt".into())
11        .with_logging()
12        .with_paged_attn(|| PagedAttentionMetaBuilder::default().build())?
13        .build()
14        .await?;
15
16    let messages = TextMessages::new()
17        .add_message(
18            TextMessageRole::System,
19            "You are an AI agent with a specialty in programming.",
20        )
21        .add_message(
22            TextMessageRole::User,
23            "Hello! How are you? Please write generic binary search function in Rust.",
24        );
25
26    let response = model.send_chat_request(messages).await?;
27
28    println!("{}", response.choices[0].message.content.as_ref().unwrap());
29    dbg!(
30        response.usage.avg_prompt_tok_per_sec,
31        response.usage.avg_compl_tok_per_sec
32    );
33
34    Ok(())
35}