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