Integrate Jinja package

Author: mattt

Sources/Hub/Config.swift

@@ -5,6 +5,7 @@
 //  Created by Piotr Kowalczuk on 06.03.25.
 
 import Foundation
+import Jinja
 
 // MARK: - Configuration files with dynamic lookup
 
@@ -413,28 +414,28 @@ public struct Config: Hashable, Sendable,
         self.dictionary(or: or)
     }
 
-    public func toJinjaCompatible() -> Any? {
+    public func jinjaValue() -> Jinja.Value {
         switch self.value {
         case let .array(val):
-            return val.map { $0.toJinjaCompatible() }
+            return .array(val.map { $0.jinjaValue() })
         case let .dictionary(val):
-            var result: [String: Any?] = [:]
+            var result: [String: Jinja.Value] = [:]
             for (key, config) in val {
-                result[key.string] = config.toJinjaCompatible()
+                result[key.string] = config.jinjaValue()
             }
-            return result
+            return .object(.init(uniqueKeysWithValues: result))
         case let .boolean(val):
-            return val
+            return .boolean(val)
         case let .floating(val):
-            return val
+            return .double(Double(String(val)) ?? Double(val))
         case let .integer(val):
-            return val
+            return .int(val)
         case let .string(val):
-            return val.string
+            return .string(val.string)
         case let .token(val):
-            return [String(val.0): val.1.string] as [String: String]
+            return [String(val.0): .string(val.1.string)]
         case .null:
-            return nil
+            return .null
         }
     }
 

Sources/Tokenizers/Tokenizer.swift

@@ -548,12 +548,12 @@ public class PreTrainedTokenizer: Tokenizer {
         }
 
         let template = try compiledTemplate(for: selectedChatTemplate)
-        var context: [String: Any] = [
-            "messages": messages,
-            "add_generation_prompt": addGenerationPrompt,
+        var context: [String: Jinja.Value] = try [
+            "messages": .array(messages.map { try Value(any: $0) }),
+            "add_generation_prompt": .boolean(addGenerationPrompt),
         ]
         if let tools {
-            context["tools"] = tools
+            context["tools"] = try .array(tools.map { try Value(any: $0) })
         }
         if let additionalContext {
             /*
@@ -562,20 +562,22 @@ public class PreTrainedTokenizer: Tokenizer {
              The default value is true in the Llama 3.1 and 3.2 chat templates, but these models will perform better if the tools are included in a system message.
              */
             for (key, value) in additionalContext {
-                context[key] = value
+                context[key] = try Value(any: value)
             }
         }
 
         for (key, value) in tokenizerConfig.dictionary(or: [:]) {
             if specialTokenAttributes.contains(key.string), !value.isNull() {
                 if let stringValue = value.string() {
-                    context[key.string] = stringValue
+                    context[key.string] = .string(stringValue)
                 } else if let dictionary = value.dictionary() {
-                    context[key.string] = addedTokenAsString(Config(dictionary))
+                    if let addedTokenString = addedTokenAsString(Config(dictionary)) {
+                        context[key.string] = .string(addedTokenString)
+                    }
                 } else if let array: [String] = value.get() {
-                    context[key.string] = array
+                    context[key.string] = .array(array.map { .string($0) })
                 } else {
-                    context[key.string] = value
+                    context[key.string] = try Value(any: value)
                 }
             }
         }

Tests/HubTests/ConfigTests.swift

@@ -431,7 +431,7 @@ class ConfigTemplatingTests: XCTestCase {
         """
 
         let got = try Template(template).render([
-            "config": cfg.toJinjaCompatible(),
+            "config": cfg.jinjaValue(),
         ])
 
         XCTAssertEqual(got, exp)