catgrad_llm/
serve.rs

1//! Abstract interfaces for serving LLMs
2use thiserror::Error;
3
4#[derive(Error, Debug)]
5pub enum Error {
6    #[error("Tokenizer Error: {0}")]
7    Tokenizer(String),
8
9    #[error("IO Error: {0}")]
10    IO(String),
11}
12
13impl From<tokenizers::tokenizer::Error> for Error {
14    fn from(err: tokenizers::tokenizer::Error) -> Self {
15        Error::Tokenizer(err.to_string())
16    }
17}
18
19/// Message type for use with instruct models
20#[derive(Clone, Debug, PartialEq, Eq)]
21pub struct Message {
22    pub role: String,
23    pub content: String,
24}
25
26pub type Result<T> = core::result::Result<T, Error>;
27
28/// A language model has a settable internal context from which it can generate new tokens.
29pub trait LM<T>: Iterator<Item = T> {
30    fn set_context(&mut self, context: Vec<T>);
31
32    fn complete(&mut self, context: Vec<T>) -> impl Iterator<Item = T> {
33        self.set_context(context);
34        self
35    }
36}
37
38/// A *loader* is conceptually a pair of language model and supporting code (tokenizers, ChatML
39/// templates, etc.)
40pub trait Loader<Token, L: LM<Token>, T: Tokenizer<Token>> {
41    fn load_runner(&self) -> Result<L>;
42    fn load_tokenizer(&self) -> Result<T>;
43}
44
45/// A [`Tokenizer`] translates between tokens and strings
46pub trait Tokenizer<Token> {
47    fn encode(&self, content: String) -> Result<Vec<Token>>;
48    fn decode(&self, tokens: Vec<Token>) -> Result<String>;
49}
50
51/// A [`Tokenizer`] which is aware of message structure
52pub trait ChatTokenizer<Token> {
53    fn encode_messages(&self, messages: Vec<Message>) -> Result<Vec<Token>>;
54}
catgrad_llm/serve.rs

catgrad_llm/
serve.rs
