Support o3/o4 models (#110)

Author: zurawiki

tiktoken-rs/README.md

@@ -105,7 +105,7 @@ println!("max_tokens: {}", max_tokens);
 
 | Encoding name           | OpenAI models                                                             |
 | ----------------------- | ------------------------------------------------------------------------- |
-| `o200k_base`            | GPT-4o models, o1 models                                                  |
+| `o200k_base`            | GPT-4o models, GPT-4.1, o1, o3, and o4 models                             |
 | `cl100k_base`           | ChatGPT models, `text-embedding-ada-002`                                  |
 | `p50k_base`             | Code models, `text-davinci-002`, `text-davinci-003`                       |
 | `p50k_edit`             | Use for edit models like `text-davinci-edit-001`, `code-davinci-edit-001` |

tiktoken-rs/src/model.rs

@@ -36,8 +36,11 @@ pub fn get_context_size(model: &str) -> usize {
         let base = rest.split(':').next().unwrap_or(rest);
         return get_context_size(base);
     }
-    if starts_with_any!(model, "o1-") {
-        return 128_000;
+    if starts_with_any!(model, "o1", "o3", "o4") {
+        return 200_000;
+    }
+    if starts_with_any!(model, "gpt-4.1") {
+        return 1_047_576;
     }
     if starts_with_any!(model, "gpt-4o") {
         return 128_000;

tiktoken-rs/src/singleton.rs

@@ -49,7 +49,7 @@ pub fn cl100k_base_singleton() -> &'static CoreBPE {
 }
 
 /// Returns a singleton instance of the o200k_base tokenizer.
-/// Use for GPT-4o models.
+/// Use for GPT-4o models and other `o` series models like `o1`, `o3`, and `o4`.
 ///
 /// This function will only initialize the tokenizer once, and then return a reference the tokenizer
 pub fn o200k_base_singleton() -> &'static CoreBPE {

tiktoken-rs/src/tiktoken_ext/openai_public.rs

@@ -118,7 +118,7 @@ pub fn cl100k_base() -> Result<CoreBPE> {
     Ok(bpe)
 }
 
-/// Use for GPT-4o models.
+/// Use for GPT-4o models and other `o` series models like `o1`, `o3`, and `o4`.
 /// Initializes and returns a new instance of the o200k_base tokenizer.
 pub fn o200k_base() -> Result<CoreBPE> {
     let o200k_base = include_str!("../../assets/o200k_base.tiktoken");

tiktoken-rs/src/tokenizer.rs

@@ -55,6 +55,7 @@ const MODEL_TO_TOKENIZER: &[(&str, Tokenizer)] = &[
     // reasoning
     ("o1", Tokenizer::O200kBase),
     ("o3", Tokenizer::O200kBase),
+    ("o4", Tokenizer::O200kBase),
     // chat
     ("gpt-4.1", Tokenizer::O200kBase),
     ("chatgpt-4o-latest", Tokenizer::O200kBase),

tiktoken-rs/src/vendor_tiktoken.rs

@@ -164,6 +164,7 @@ impl std::fmt::Display for DecodeKeyError {
 impl std::error::Error for DecodeKeyError {}
 
 #[derive(Debug, Clone)]
+#[allow(dead_code)]
 pub struct DecodeError {
     pub message: String,
 }

tiktoken-rs/tests/model.rs

@@ -11,3 +11,9 @@ fn test_finetuned_context_size() {
         get_context_size("gpt-4o")
     );
 }
+
+#[test]
+fn test_o_series_context_size() {
+    assert_eq!(get_context_size("o3-small"), 200_000);
+    assert_eq!(get_context_size("o4"), 200_000);
+}