kalosm-language-model 0.1.0

A common interface for language models/transformers
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

pub use crate::local::session::*;
use crate::{embedding::Embedding, model::*};
use kalosm_sample::Tokenizer;
use kalosm_streams::ChannelTextStream;
use llm::InferenceSessionConfig;
use llm_samplers::prelude::Sampler;
use std::sync::Arc;
use std::sync::Mutex;

mod download;
mod session;

pub(crate) trait LocalModelType {
    fn model_type() -> ModelType;
}

macro_rules! local_model {
    ($ty: expr, $space: ident) => {
        impl LocalModelType for LocalSession<$space> {
            fn model_type() -> ModelType {
                $ty
            }
        }

        #[async_trait::async_trait]
        impl crate::model::CreateModel for LocalSession<$space> {
            async fn start() -> Self {
                let model = Self::model_type().download().await;
                let session = model.start_session(InferenceSessionConfig {
                    n_batch: 64,
                    n_threads: num_cpus::get(),
                    ..Default::default()
                });

                LocalSession::new(model, session)
            }

            fn requires_download() -> bool {
                Self::model_type().requires_download()
            }
        }

        #[async_trait::async_trait]
        impl crate::model::Model for LocalSession<$space> {
            type TextStream = ChannelTextStream<String>;
            type SyncModel = crate::SyncModelNotSupported;

            fn tokenizer(&self) -> Arc<dyn Tokenizer + Send + Sync> {
                self.get_tokenizer() as Arc<dyn Tokenizer + Send + Sync>
            }

            async fn stream_text_inner(
                &self,
                prompt: &str,
                generation_parameters: GenerationParameters,
            ) -> anyhow::Result<Self::TextStream> {
                Ok(self.infer(prompt.to_string(), generation_parameters).await)
            }

            async fn stream_text_with_sampler(
                &self,
                prompt: &str,
                max_tokens: Option<u32>,
                stop_on: Option<&str>,
                sampler: Arc<Mutex<dyn Sampler>>,
            ) -> anyhow::Result<Self::TextStream> {
                Ok(self
                    .infer_sampler(prompt.to_string(), max_tokens, stop_on, sampler)
                    .await)
            }
        }

        #[async_trait::async_trait]
        impl crate::model::Embedder<$space> for LocalSession<$space> {
            async fn embed(&mut self, input: &str) -> anyhow::Result<Embedding<$space>> {
                self.get_embedding(input).await
            }

            async fn embed_batch(
                &mut self,
                inputs: &[&str],
            ) -> anyhow::Result<Vec<Embedding<$space>>> {
                let mut result = Vec::new();
                for input in inputs {
                    result.push(self.get_embedding(input).await?);
                }
                Ok(result)
            }
        }
    };
}

local_model!(ModelType::Llama(LlamaType::Vicuna), VicunaSpace);
local_model!(ModelType::Llama(LlamaType::Guanaco), GuanacoSpace);
local_model!(ModelType::Llama(LlamaType::WizardLm), WizardLmSpace);
local_model!(ModelType::Llama(LlamaType::Orca), OrcaSpace);
local_model!(
    ModelType::Llama(LlamaType::LlamaSevenChat),
    LlamaSevenChatSpace
);
local_model!(
    ModelType::Llama(LlamaType::LlamaThirteenChat),
    LlamaThirteenChatSpace
);
local_model!(ModelType::Mpt(MptType::Base), MptBaseSpace);
local_model!(ModelType::Mpt(MptType::Story), MptStorySpace);
local_model!(ModelType::Mpt(MptType::Instruct), MptInstructSpace);
local_model!(ModelType::Mpt(MptType::Chat), MptChatSpace);
local_model!(
    ModelType::GptNeoX(GptNeoXType::LargePythia),
    LargePythiaSpace
);
local_model!(ModelType::GptNeoX(GptNeoXType::TinyPythia), TinyPythiaSpace);
local_model!(
    ModelType::GptNeoX(GptNeoXType::DollySevenB),
    DollySevenBSpace
);
local_model!(ModelType::GptNeoX(GptNeoXType::StableLm), StableLmSpace);

pub(crate) fn get_embeddings<S: crate::VectorSpace>(
    model: &dyn llm::Model,
    embed: &str,
) -> Embedding<S> {
    let mut session = model.start_session(Default::default());
    let mut output_request = llm::OutputRequest {
        all_logits: None,
        embeddings: Some(Vec::new()),
    };
    let _ = session.feed_prompt(model, embed, &mut output_request, |_| {
        Ok::<_, std::convert::Infallible>(llm::InferenceFeedback::Halt)
    });
    Embedding::from(output_request.embeddings.unwrap())
}