erniebot-rs

非官方百度 Ernie(Wenxin Yiyan, Qianfan) Rust SDK，目前支持三个模块：聊天、文本嵌入（embedding）和文本生成图像（text2image）。

2024/04/09 更新：添加对 bce-reranker-base-v1 重新排序模型的支持

2024/04/21 更新 对于同步模式，使用 ureq 替代 reqwest_blocking，因此可以提高与 tokio 的兼容性。

安装

将以下内容添加到您的 Cargo.toml 文件中

[dependencies]
erniebot-rs = "0.4.1"

身份验证

在使用之前，将 AK 和 SK 导入环境变量

export QIANFAN_AK=***  
export QIANFAN_SK=***

聊天

默认支持的模型包括

• ErnieBotTurbo
• ErnieBot
• Ernie40

以下方式调用这些模型

fn test_invoke() {  
    let chat = ChatEndpoint::new(ChatModel::ErnieBotTurbo).unwrap();  
    let messages = vec![  
        Message {  
            role: Role::User,  
            content: "hello, I'm a developer. I'm developing a rust SDK for qianfan LLM. If you get this message, that means I successfully send you this message using invoke method".to_string(),  
            ..Default::default()  
        },  
    ];  
    let options = vec![  
        ChatOpt::Temperature(0.5),  
        ChatOpt::TopP(0.5),  
        ChatOpt::TopK(50),  
    ];  
    let response = chat.invoke(&messages, &options).unwrap();  
    let result = response.get_chat_result().unwrap();  
    println!("{}", result);  
}

对于其他模型，使用 new_with_custom_endpoint 调用。字段名是 Qianfan API 的最后一部分。以 llama_2_70b 为例，API 地址是 https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/llama_2_70b。调用方法如下

fn test_custom_endpoint() {  
    let chat = ChatEndpoint::new_with_custom_endpoint("llama_2_70b").unwrap();  
    let messages = vec![  
        Message {  
            role: Role::User,  
            content: "hello, I'm a developer. I'm developing a rust SDK for qianfan LLM. If you get this message, that means I successfully send you this message using a custom endpoint".to_string(),  
            ..Default::default()  
        },  
    ];  
    let options = Vec::new();  
    let response = chat.invoke(&messages, &options).unwrap();  
    let result = response.get_chat_result().unwrap();  
    println!("{}", result);  
}

支持四种调用方法：invoke（同步非流式）、ainvoke（异步非流式）、stream（同步流式）和 astream（异步流式）。

例如，astream 调用方法如下

fn test_astream() {  
    let chat = ChatEndpoint::new(ChatModel::ErnieBotTurbo).unwrap();  
    let messages = vec![  
        Message {  
            role: Role::User,  
            content: "hello, I'm a developer. I'm developing a rust SDK for qianfan LLM. If you get this message, that means I successfully send you this message using async stream method. Now reply to me a message as long as possible so that I can test if this function doing well".to_string(),  
            ..Default::default()  
        },  
    ];  
    let options = Vec::new();  
    let rt = Runtime::new().unwrap();  
    rt.block_on(async move {  
        let mut stream_response = chat.astream(&messages, &options).await.unwrap();  
        while let Some(response) = stream_response.next().await {  
            let result = response.get_chat_result().unwrap();  
            print!("{}", result);  
            //flush  
            std::io::stdout().flush().unwrap();  
        }  
    });  
    println!();  
}

对于某些模型，如 ErnieBot，它们支持传递函数进行调用的选项。您可以参考 examples/chat_with_function.rs 中的示例。

请注意，由于每个特定模型的参数要求不同，此 SDK 不执行本地参数验证，而是将参数传递到服务器进行验证。因此，如果参数不符合要求，服务器将返回错误消息。

嵌入

支持目前（截至 2024/02/26）Qianfan 平台上可用的四个模型

• EmbeddingV1
• BgeLargeZh
• BgeLargeEn
• Tao8k

调用方式与聊天类似，支持 invoke 和 ainvoke 模式

fn test_async_embedding() {  
    let embedding = EmbeddingEndpoint::new(EmbeddingModel::EmbeddingV1).unwrap();  
    let input = vec![  
        "你好".to_string(),  
        "你叫什么名字".to_string(),  
        "你是谁".to_string(),  
    ];  
    let rt = Runtime::new().unwrap();  
    let embedding_response = rt.block_on(embedding.ainvoke(&input, None)).unwrap();  
    let embedding_results = embedding_response.get_embedding_results().unwrap();  
    println!("{},{}", embedding_results.len(), embedding_results[0].len());  
}

文本生成图像

支持默认的 StableDiffusionXL，并允许使用自定义模型（如 Wenxin Yige）。

fn main() {
    let text2image = Text2ImageEndpoint::new(Text2ImageModel::StableDiffusionXL).unwrap();
    let prompt = "A beautiful sunset over the ocean".to_string();
    let options = vec![
        Text2ImageOpt::Style(Style::DigitalArt),
        Text2ImageOpt::Size(Size::S1024x768),
    ];
    let text2image_response = text2image.invoke(&prompt, &options).unwrap();
    let image_results = text2image_response.get_image_results().unwrap();
    for (index, image_string) in image_results.into_iter().enumerate() {
        let image = base64_to_image(image_string).unwrap();
        let filepath = format!("./tmp/image_{}.png", index);
        image.save(filepath).unwrap();
    }
}

待办事项

• 文档
• 聊天模型众多，因此选项可能不完整，需要进一步补充。
• 需要更多测试。
• 福裕-8B图像理解模型尚不支持。