another_tiktoken_rs/
model.rs

1/*!
2 * contains information about OpenAI models.
3 */
4
5/// Macro to check if a given str starts with any of the specified prefixes.
6macro_rules! starts_with_any {
7    ($str:expr, $($prefix:expr),* $(,)?) => {
8        false $(|| $str.starts_with($prefix))*
9    };
10}
11
12/// Returns the context size of a specified model.
13///
14/// The context size represents the maximum number of tokens a model can process in a single input.
15/// This function checks the model name and returns the corresponding context size.
16/// See <https://platform.openai.com/docs/models> for up-to-date information.
17///
18/// # Arguments
19///
20/// * `model` - A string slice that holds the name of the model.
21///
22/// # Examples
23///
24/// ```
25/// use another_tiktoken_rs::model::get_context_size;
26/// let model = "gpt-4-32k";
27/// let context_size = get_context_size(model);
28/// assert_eq!(context_size, 32768);
29/// ```
30///
31/// # Panics
32///
33/// This function does not panic. It returns a default value of 4096 if the model is not recognized.
34pub fn get_context_size(model: &str) -> usize {
35    if starts_with_any!(model, "gpt-4-32k") {
36        return 32_768;
37    }
38    if starts_with_any!(model, "gpt-4") {
39        return 8192;
40    }
41    if starts_with_any!(model, "gpt-3.5-turbo-16k") {
42        return 16_384;
43    }
44    if starts_with_any!(model, "gpt-3.5-turbo") {
45        return 4096;
46    }
47    if starts_with_any!(model, "text-davinci-002", "text-davinci-003") {
48        return 4097;
49    }
50    if starts_with_any!(model, "ada", "babbage", "curie") {
51        return 2049;
52    }
53    if starts_with_any!(model, "code-cushman-001") {
54        return 2048;
55    }
56    if starts_with_any!(model, "code-davinci-002") {
57        return 8001;
58    }
59    if starts_with_any!(model, "davinci") {
60        return 2049;
61    }
62    if starts_with_any!(model, "text-ada-001", "text-babbage-001", "text-curie-001") {
63        return 2049;
64    }
65    4096
66}
another_tiktoken_rs/model.rs

another_tiktoken_rs/
model.rs
