Add model

Author: aespinilla

.gitignore

@@ -4,6 +4,5 @@
 /*.xcodeproj
 xcuserdata/
 DerivedData/
-.swiftpm/config/registries.json
-.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata
+.swiftpm
 .netrc

Sources/Tiktoken/Model.swift

@@ -0,0 +1,74 @@
+//
+//  Model.swift
+//  
+//
+//  Created by Alberto Espinilla Garrido on 20/3/23.
+//
+
+import Foundation
+
+enum Model {
+    static func getEncoding(_ name: String) -> String? {
+        if let encodingName = MODEL_TO_ENCODING[name] {
+            return encodingName
+        }
+        return findPrefix(with: name)
+    }
+}
+
+private extension Model {
+    static let MODEL_PREFIX_TO_ENCODING: [String: String] = [
+        // chat
+        "gpt-4-": "cl100k_base",  // e.g., gpt-4-0314, etc., plus gpt-4-32k
+        "gpt-3.5-turbo-": "cl100k_base",  // e.g, gpt-3.5-turbo-0301, -0401, etc.
+    ]
+    
+    static let MODEL_TO_ENCODING: [String: String] = [
+        // chat
+        "gpt-4": "cl100k_base",
+        "gpt-3.5-turbo": "cl100k_base",
+        // text
+        "text-davinci-003": "p50k_base",
+        "text-davinci-002": "p50k_base",
+        "text-davinci-001": "r50k_base",
+        "text-curie-001": "r50k_base",
+        "text-babbage-001": "r50k_base",
+        "text-ada-001": "r50k_base",
+        "davinci": "r50k_base",
+        "curie": "r50k_base",
+        "babbage": "r50k_base",
+        "ada": "r50k_base",
+        // code
+        "code-davinci-002": "p50k_base",
+        "code-davinci-001": "p50k_base",
+        "code-cushman-002": "p50k_base",
+        "code-cushman-001": "p50k_base",
+        "davinci-codex": "p50k_base",
+        "cushman-codex": "p50k_base",
+        // edit
+        "text-davinci-edit-001": "p50k_edit",
+        "code-davinci-edit-001": "p50k_edit",
+        // embeddings
+        "text-embedding-ada-002": "cl100k_base",
+        // old embeddings
+        "text-similarity-davinci-001": "r50k_base",
+        "text-similarity-curie-001": "r50k_base",
+        "text-similarity-babbage-001": "r50k_base",
+        "text-similarity-ada-001": "r50k_base",
+        "text-search-davinci-doc-001": "r50k_base",
+        "text-search-curie-doc-001": "r50k_base",
+        "text-search-babbage-doc-001": "r50k_base",
+        "text-search-ada-doc-001": "r50k_base",
+        "code-search-babbage-code-001": "r50k_base",
+        "code-search-ada-code-001": "r50k_base",
+        // open source
+        "gpt2": "gpt2",
+    ]
+    
+    static func findPrefix(with name: String) -> String? {
+        guard let key = Model.MODEL_PREFIX_TO_ENCODING.keys.first(where: { name.starts(with: $0) }) else {
+            return nil
+        }
+        return Model.MODEL_PREFIX_TO_ENCODING[key]
+    }
+}

Sources/Tiktoken/Tiktoken.swift

@@ -3,4 +3,18 @@ public struct Tiktoken {
 
     public init() {
     }
+    
+    public func getEncoding(_ name: String) -> Encoding? {
+        nil
+    }
+    
+    public func getEncoding(for model: String) -> Encoding? {
+        nil
+    }
+}
+
+
+public protocol Encoding {
+    func encode(value: String) -> [Int]
+    func decode(value: [Int]) -> String
 }

Tests/TiktokenTests/ModelTests.swift

@@ -0,0 +1,43 @@
+//
+//  ModelTests.swift
+//  
+//
+//  Created by Alberto Espinilla Garrido on 20/3/23.
+//
+
+import XCTest
+@testable import Tiktoken
+
+final class ModelTests: XCTestCase {
+
+    func testGivenModelNamesWhenGetEncodingThenMatch() throws {
+        try [
+            Test(input: "gpt-4", output: "cl100k_base"),
+            Test(input: "gpt-3.5-turbo", output: "cl100k_base"),
+            Test(input: "davinci", output: "r50k_base"),
+            Test(input: "text-davinci-edit-001", output: "p50k_edit"),
+        ].forEach({
+            let output = Model.getEncoding($0.input)
+            XCTAssertEqual(try XCTUnwrap(output), $0.output)
+        })
+    }
+    
+    func testGivenModelNamesWithPrefisWhenGetEncodingThenMatch() throws {
+        try [
+            Test(input: "gpt-4-0314", output: "cl100k_base"),
+            Test(input: "gpt-4-32k", output: "cl100k_base"),
+            Test(input: "gpt-3.5-turbo-0301", output: "cl100k_base"),
+            Test(input: "gpt-3.5-turbo-0401", output: "cl100k_base"),
+        ].forEach({
+            let output = Model.getEncoding($0.input)
+            XCTAssertEqual(try XCTUnwrap(output), $0.output)
+        })
+    }
+    
+    func testGivenUnknowModelNamesWhenGetEncodingThenMatchNil() throws {
+        ["sample", "chatgpt", "invalid", "test"].forEach({
+            let output = Model.getEncoding($0)
+            XCTAssertNil(output)
+        })
+    }
+}

Tests/TiktokenTests/Test.swift

@@ -0,0 +1,13 @@
+//
+//  Test.swift
+//  
+//
+//  Created by Alberto Espinilla Garrido on 20/3/23.
+//
+
+import Foundation
+
+struct Test<Input, Output> {
+    let input: Input
+    let output: Output
+}