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 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}