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