feat: Support Starcoder2 (#20)

* feat: Support Starcoder2
2024-03-08 13:28:37 +08:00
parent e876e18605
commit a94bf79d7e
4 changed files with 281 additions and 2 deletions
--- a/Libraries/LLM/Configuration.swift
+++ b/Libraries/LLM/Configuration.swift
@@ -32,6 +32,7 @@ public enum ModelType: String, Codable {
    case phi
    case gemma
    case qwen2
+    case starcoder2

    func createModel(configuration: URL) throws -> LLMModel {
        switch self {
@@ -51,6 +52,10 @@ public enum ModelType: String, Codable {
            let configuration = try JSONDecoder().decode(
                Qwen2Configuration.self, from: Data(contentsOf: configuration))
            return Qwen2Model(configuration)
+        case .starcoder2:
+            let configuration = try JSONDecoder().decode(
+                Starcoder2Configuration.self, from: Data(contentsOf: configuration))
+            return Starcoder2Model(configuration)
        }
    }
 }