llama-crab 0.1.8

Safe, ergonomic and complete Rust bindings for llama.cpp
Documentation 
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

//! Direct high-level convenience helpers.

use crate::{chat::ChatMessage, error::Result};

use super::{CompletionOptions, Llama};

/// Generated text plus token accounting for a single prompt.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TextGeneration {
    /// Generated text.
    pub text: String,
    /// Number of tokens in the prompt.
    pub prompt_tokens: u32,
    /// Number of tokens generated by the model.
    pub completion_tokens: u32,
    /// Prompt plus generated token count.
    pub total_tokens: u32,
}

/// Generated assistant message plus token accounting.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChatGeneration {
    /// Assistant message returned by the model.
    pub message: ChatMessage,
    /// Approximate token count for the input messages.
    pub prompt_tokens: u32,
    /// Number of tokens in the assistant message.
    pub completion_tokens: u32,
    /// Prompt plus generated token count.
    pub total_tokens: u32,
}

/// Embedding vectors plus token accounting.
#[derive(Debug, Clone, PartialEq)]
pub struct EmbeddingBatch {
    /// One embedding vector per input text.
    pub vectors: Vec<Vec<f32>>,
    /// Number of tokens across all input texts.
    pub prompt_tokens: u32,
    /// Same value as `prompt_tokens`; embeddings do not generate output tokens.
    pub total_tokens: u32,
}

impl Llama {
    /// Generate text and return token counts alongside the decoded output.
    pub fn generate_text(&mut self, prompt: &str, max_tokens: usize) -> Result<TextGeneration> {
        let prompt_tokens = self.model().tokenize(prompt, true, true)?.len() as u32;
        let completion =
            self.create_completion_with_options(prompt, CompletionOptions::new(max_tokens))?;
        let completion_tokens = completion.n_tokens as u32;

        Ok(TextGeneration {
            text: completion.text,
            prompt_tokens,
            completion_tokens,
            total_tokens: prompt_tokens + completion_tokens,
        })
    }

    /// Generate a chat response and return token counts with the assistant
    /// message.
    pub fn generate_chat(
        &mut self,
        messages: &[ChatMessage],
        max_tokens: usize,
    ) -> Result<ChatGeneration> {
        let prompt_tokens = messages
            .iter()
            .map(|message| {
                self.model()
                    .tokenize(&message.content, true, true)
                    .map_or(0, |tokens| tokens.len() as u32)
            })
            .sum::<u32>();
        let message = self.create_chat_completion(messages, max_tokens)?;
        let completion_tokens = self.model().tokenize(&message.content, false, true)?.len() as u32;

        Ok(ChatGeneration {
            message,
            prompt_tokens,
            completion_tokens,
            total_tokens: prompt_tokens + completion_tokens,
        })
    }

    /// Embed many texts and return token counts for the input batch.
    pub fn embed_texts(&mut self, texts: &[String], normalize: bool) -> Result<EmbeddingBatch> {
        let mut prompt_tokens = 0_u32;
        let mut vectors = Vec::with_capacity(texts.len());

        for text in texts {
            prompt_tokens += self.model().tokenize(text, true, false)?.len() as u32;
            vectors.push(self.embed(text, normalize)?);
        }

        Ok(EmbeddingBatch {
            vectors,
            prompt_tokens,
            total_tokens: prompt_tokens,
        })
    }
}