Merge pull request #104 from neonredwood/support-gpt4-omni

Add support for GPT-4-O, "Omni" model

Author: dqbd

js/src/core.ts

@@ -274,6 +274,10 @@ export function getEncodingNameForModel(model: TiktokenModel) {
     case "text-embedding-ada-002": {
       return "cl100k_base";
     }
+    case "gpt-4o":
+    case "gpt-4o-2024-05-13": {
+      return "o200k_base";
+    }
     default:
       never(model);
       throw new Error("Unknown model");

js/src/index.ts

@@ -4,6 +4,7 @@ import p50k_base from "./ranks/p50k_base";
 import p50k_edit from "./ranks/p50k_edit";
 import r50k_base from "./ranks/r50k_base";
 import cl100k_base from "./ranks/cl100k_base";
+import o200k_base from "./ranks/o200k_base";
 
 import { Tiktoken, getEncodingNameForModel } from "./core";
 import { never } from "./utils";
@@ -23,6 +24,8 @@ export function getEncoding(
       return new Tiktoken(p50k_edit, extendSpecialTokens);
     case "cl100k_base":
       return new Tiktoken(cl100k_base, extendSpecialTokens);
+    case "o200k_base":
+      return new Tiktoken(o200k_base, extendSpecialTokens);
     default:
       never(encoding);
       throw new Error("Unknown encoding");

scripts/ranks.ts

@@ -264,8 +264,8 @@ async function main() {
       const bpe = await downloadBpe(data);
 
       if (lib === "js") {
-        bpe.pat_str = bpe.pat_str.replace(
-          /\(\?i:(.*?)\)/,
+        bpe.pat_str = bpe.pat_str.replaceAll(
+          /\(\?i:(.*?)\)/g,
           (_, match: string) =>
             `(${match
               .split("|")

tiktoken/model_to_encoding.json

@@ -52,5 +52,7 @@
     "gpt-4-turbo-preview": "cl100k_base",
     "gpt-4-1106-preview": "cl100k_base",
     "gpt-4-0125-preview": "cl100k_base",
-    "gpt-4-vision-preview": "cl100k_base"
+    "gpt-4-vision-preview": "cl100k_base",
+    "gpt-4o": "o200k_base",
+    "gpt-4o-2024-05-13": "o200k_base"
 }

tiktoken/registry.json

@@ -46,5 +46,13 @@
             "<|endofprompt|>": 100276
         },
         "pat_str": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
+    },
+    "o200k_base": {
+      "load_tiktoken_bpe": "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken",
+      "special_tokens": {
+        "<|endoftext|>": 199999,
+        "<|endofprompt|>": 200018
+      },
+      "pat_str": "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
     }
-}
+}

wasm/src/lib.rs

@@ -133,6 +133,19 @@ impl CoreBPEConstructor {
             include_str!("./ranks/cl100k_base.regex.tiktoken"),
         )
     }
+
+    #[cfg(feature = "inline")]
+    fn o200k_base() -> Self {
+        let mut special_tokens = HashMap::default();
+        special_tokens.insert(String::from(ENDOFTEXT), 199999);
+        special_tokens.insert(String::from(ENDOFPROMPT), 200018);
+
+        CoreBPEConstructor::new(
+            include_str!("./ranks/o200k_base.compress.tiktoken"),
+            Some(special_tokens),
+            include_str!("./ranks/o200k_base.regex.tiktoken"),
+        )
+    }
 }
 
 #[wasm_bindgen]
@@ -179,6 +192,7 @@ impl Tiktoken {
             "p50k_base" => Ok(CoreBPEConstructor::p50k_base()),
             "p50k_edit" => Ok(CoreBPEConstructor::p50k_edit()),
             "cl100k_base" => Ok(CoreBPEConstructor::cl100k_base()),
+            "o200k_base" => Ok(CoreBPEConstructor::o200k_base()),
             &_ => Err(JsError::new("Invalid encoding")),
         }?;
 
@@ -325,7 +339,7 @@ impl Tiktoken {
 #[cfg(feature = "inline")]
 #[wasm_bindgen(typescript_custom_section)]
 const _: &'static str = r#"
-export type TiktokenEncoding = "gpt2" | "r50k_base" | "p50k_base" | "p50k_edit" | "cl100k_base"; 
+export type TiktokenEncoding = "gpt2" | "r50k_base" | "p50k_base" | "p50k_edit" | "cl100k_base" | "o200k_base";
 
 /**
  * @param {TiktokenEncoding} encoding
@@ -404,6 +418,8 @@ export type TiktokenModel =
     | "gpt-4-1106-preview"
     | "gpt-4-0125-preview"
     | "gpt-4-vision-preview"
+    | "gpt-4o"
+    | "gpt-4o-2024-05-13"
 
 /**
  * @param {TiktokenModel} encoding
@@ -474,6 +490,8 @@ pub fn encoding_for_model(
         "gpt-4-turbo-2024-04-09" => Ok("cl100k_base"),
         "gpt-4-turbo-preview" => Ok("cl100k_base"),
         "gpt-4-0125-preview" => Ok("cl100k_base"),
+        "gpt-4o" => Ok("o200k_base"),
+        "gpt-4o-2024-05-13" => Ok("o200k_base"),
         model => Err(JsError::new(
             format!("Invalid model: {}", model.to_string()).as_str(),
         )),

wasm/test/test_simple_public.test.ts

@@ -59,6 +59,28 @@ describe("cl100k_base", () => {
   });
 });
 
+describe("o200k_base", () => {
+  const enc = get_encoding("o200k_base");
+
+  it("encodes hello world string", () => {
+    expect(enc.encode("hello world")).toStrictEqual(
+      new Uint32Array([24912, 2375])
+    );
+  });
+
+  it("decodes hello world string", () => {
+    expect(
+      new TextDecoder().decode(enc.decode(new Uint32Array([24912, 2375])))
+    ).toStrictEqual("hello world");
+  });
+
+  it("encodes hello world string, all allowed special characters", () => {
+    expect(enc.encode("hello <|endoftext|>", "all")).toStrictEqual(
+      new Uint32Array([24912, 220, 199999])
+    );
+  });
+});
+
 it("test_simple", () => {
   const encodings = [
     "gpt2",