-
Notifications
You must be signed in to change notification settings - Fork 0
/
demo_bedrock_raw_mistral.rs
85 lines (67 loc) · 2.28 KB
/
demo_bedrock_raw_mistral.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
use futures::TryStreamExt;
use std::io;
use std::io::Write;
use crate::bedrock::bedrock_client::{BedrockClient, BedrockClientOptions};
pub async fn demo_bedrock_mistral_raw_stream(model_id: &str, prompt: &str) {
let profile_name = "bedrock";
let region = "us-west-2";
let payload = serde_json::json!({
"prompt": prompt,
"max_tokens" : 200,
"stop" : ["[INST]"],
"temperature": 0.5,
"top_p": 0.9,
"top_k": 100,
});
let options = BedrockClientOptions::new()
.profile_name(profile_name)
.region(region);
let client = BedrockClient::new(options).await;
let stream = client
.generate_raw_stream(
model_id.to_string(),
payload,
)
.await;
let stream = match stream {
Ok(stream) => stream,
Err(err) => {
println!("Error: {:?}", err);
return;
}
};
// consumme the stream and print the response
stream
.try_for_each(|chunk| async move {
println!("{:?}", chunk);
// Flush the output to ensure the prompt is displayed.
io::stdout().flush().unwrap();
Ok(())
})
.await
.unwrap();
}
// Write a test
#[cfg(test)]
mod tests {
use super::*;
use crate::bedrock::model_info::{ModelInfo, ModelName};
#[tokio::test]
async fn test_demo_bedrock_mistral_raw_stream_8x7() {
let model_id = ModelInfo::from_model_name(ModelName::MistralMixtral8X7BInstruct0x);
let prompt = "<s>[INST] What is your favourite condiment? [/INST]";
demo_bedrock_mistral_raw_stream(&model_id,&prompt).await;
}
#[tokio::test]
async fn test_demo_bedrock_mistral_raw_stream_7b() {
let model_id = ModelInfo::from_model_name(ModelName::MistralMistral7BInstruct0x);
let prompt = "<s>[INST] What is your favourite condiment? [/INST]";
demo_bedrock_mistral_raw_stream(&model_id,&prompt).await;
}
#[tokio::test]
async fn test_demo_bedrock_mistral_raw_stream_large() {
let model_id = ModelInfo::from_model_name(ModelName::MistralLarge);
let prompt = "<s>[INST] What is your favourite condiment? [/INST]";
demo_bedrock_mistral_raw_stream(&model_id,&prompt).await;
}
}