Add gpt-4-0125 and gpt-3.5-turbo-1106 context sizes (#59)

zurawiki · Jan 26, 2024 · f84907c · f84907c
1 parent b220529
commit f84907c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/tiktoken-rs/src/model.rs b/tiktoken-rs/src/model.rs
@@ -32,6 +32,9 @@ macro_rules! starts_with_any {
 ///
 /// This function does not panic. It returns a default value of 4096 if the model is not recognized.
 pub fn get_context_size(model: &str) -> usize {
+    if starts_with_any!(model, "gpt-4-0125") {
+        return 128_000;
+    }
     if starts_with_any!(model, "gpt-4-1106") {
         return 128_000;
     }
@@ -41,6 +44,9 @@ pub fn get_context_size(model: &str) -> usize {
     if starts_with_any!(model, "gpt-4") {
         return 8192;
     }
+    if starts_with_any!(model, "gpt-3.5-turbo-1106") {
+        return 16_385;
+    }
     if starts_with_any!(model, "gpt-3.5-turbo-16k") {
         return 16_384;
     }

diff --git a/tiktoken-rs/src/tokenizer.rs b/tiktoken-rs/src/tokenizer.rs
@@ -136,11 +136,19 @@ mod tests {
 
     #[test]
     fn test_get_tokenizer() {
+        assert_eq!(
+            get_tokenizer("gpt-4-0125-preview"),
+            Some(Tokenizer::Cl100kBase)
+        );
         assert_eq!(get_tokenizer("gpt-4-32k-0314"), Some(Tokenizer::Cl100kBase));
         assert_eq!(
             get_tokenizer("gpt-4-1106-preview"),
             Some(Tokenizer::Cl100kBase)
         );
+        assert_eq!(
+            get_tokenizer("gpt-3.5-turbo-1106"),
+            Some(Tokenizer::Cl100kBase),
+        );
         assert_eq!(get_tokenizer("gpt-3.5-turbo"), Some(Tokenizer::Cl100kBase));
         assert_eq!(
             get_tokenizer("ft:gpt-3.5-turbo:XXXXXX:2023-11-11"),