spin_sdk/
llm.rs

1pub use crate::wit::v2::llm::{
2    self, EmbeddingsResult, EmbeddingsUsage, Error, InferencingParams, InferencingResult,
3    InferencingUsage,
4};
5
6/// The model use for inferencing
7#[allow(missing_docs)]
8#[derive(Debug, Clone, Copy)]
9pub enum InferencingModel<'a> {
10    Llama2Chat,
11    CodellamaInstruct,
12    Other(&'a str),
13}
14
15impl<'a> std::fmt::Display for InferencingModel<'a> {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        let str = match self {
18            InferencingModel::Llama2Chat => "llama2-chat",
19            InferencingModel::CodellamaInstruct => "codellama-instruct",
20            InferencingModel::Other(s) => s,
21        };
22        f.write_str(str)
23    }
24}
25
26impl Default for InferencingParams {
27    fn default() -> Self {
28        Self {
29            max_tokens: 100,
30            repeat_penalty: 1.1,
31            repeat_penalty_last_n_token_count: 64,
32            temperature: 0.8,
33            top_k: 40,
34            top_p: 0.9,
35        }
36    }
37}
38
39/// Perform inferencing using the provided model and prompt
40pub fn infer(model: InferencingModel, prompt: &str) -> Result<InferencingResult, Error> {
41    llm::infer(&model.to_string(), prompt, None)
42}
43
44/// Perform inferencing using the provided model, prompt, and options
45pub fn infer_with_options(
46    model: InferencingModel,
47    prompt: &str,
48    options: InferencingParams,
49) -> Result<InferencingResult, Error> {
50    llm::infer(&model.to_string(), prompt, Some(options))
51}
52
53/// Model used for generating embeddings
54#[allow(missing_docs)]
55#[derive(Debug, Clone, Copy)]
56pub enum EmbeddingModel<'a> {
57    AllMiniLmL6V2,
58    Other(&'a str),
59}
60
61impl<'a> std::fmt::Display for EmbeddingModel<'a> {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        let str = match self {
64            EmbeddingModel::AllMiniLmL6V2 => "all-minilm-l6-v2",
65            EmbeddingModel::Other(s) => s,
66        };
67        f.write_str(str)
68    }
69}
70
71/// Generate embeddings using the provided model and collection of text
72pub fn generate_embeddings(
73    model: EmbeddingModel,
74    text: &[String],
75) -> Result<llm::EmbeddingsResult, Error> {
76    llm::generate_embeddings(&model.to_string(), text)
77}
spin_sdk/llm.rs

spin_sdk/
llm.rs
