Skip to main content

simple/
main.rs

1use anyhow::Result;
2use mistralrs::{
3    IsqType, PagedAttentionMetaBuilder, RequestBuilder, TextMessageRole, TextMessages,
4    TextModelBuilder,
5};
6
7#[tokio::main]
8async fn main() -> Result<()> {
9    let model = TextModelBuilder::new("microsoft/Phi-3.5-mini-instruct")
10        .with_isq(IsqType::Q8_0)
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    // Next example: Return some logprobs with the `RequestBuilder`, which enables higher configurability.
35    let request = RequestBuilder::new().return_logprobs(true).add_message(
36        TextMessageRole::User,
37        "Please write a mathematical equation where a few numbers are added.",
38    );
39
40    let response = model.send_chat_request(request).await?;
41
42    println!(
43        "Logprobs: {:?}",
44        &response.choices[0]
45            .logprobs
46            .as_ref()
47            .unwrap()
48            .content
49            .as_ref()
50            .unwrap()[0..3]
51    );
52
53    Ok(())
54}