Skip to main content

async/
main.rs

1use anyhow::Result;
2use mistralrs::{
3    IsqType, PagedAttentionMetaBuilder, RequestBuilder, TextMessageRole, TextMessages,
4    TextModelBuilder,
5};
6
7async fn run() -> Result<()> {
8    let model = TextModelBuilder::new("microsoft/Phi-3.5-mini-instruct")
9        .with_isq(IsqType::Q8_0)
10        .with_logging()
11        .with_paged_attn(|| PagedAttentionMetaBuilder::default().build())?
12        .build()
13        .await?;
14
15    let messages = TextMessages::new()
16        .add_message(
17            TextMessageRole::System,
18            "You are an AI agent with a specialty in programming.",
19        )
20        .add_message(
21            TextMessageRole::User,
22            "Hello! How are you? Please write generic binary search function in Rust.",
23        );
24
25    let response = model.send_chat_request(messages).await?;
26
27    println!("{}", response.choices[0].message.content.as_ref().unwrap());
28    dbg!(
29        response.usage.avg_prompt_tok_per_sec,
30        response.usage.avg_compl_tok_per_sec
31    );
32
33    // Next example: Return some logprobs with the `RequestBuilder`, which enables higher configurability.
34    let request = RequestBuilder::new().return_logprobs(true).add_message(
35        TextMessageRole::User,
36        "Please write a mathematical equation where a few numbers are added.",
37    );
38
39    let response = model.send_chat_request(request).await?;
40
41    println!(
42        "Logprobs: {:?}",
43        &response.choices[0]
44            .logprobs
45            .as_ref()
46            .unwrap()
47            .content
48            .as_ref()
49            .unwrap()[0..3]
50    );
51    Ok(())
52}
53
54#[tokio::main]
55async fn main() -> Result<()> {
56    let task = tokio::task::spawn(async move { run().await.unwrap() });
57    task.await?;
58    Ok(())
59}