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

@@ -549,32 +549,34 @@ 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 {
             // Additional keys and values to be added to the context provided to the prompt templating engine.
             // For example, the app could set "tools_in_user_message" to false for Llama 3.1 and 3.2 if a system message is provided.
             // 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

@@ -435,7 +435,7 @@ struct ConfigTests {
             """
 
         let got = try Template(template).render([
-            "config": cfg.toJinjaCompatible()
+            "config": cfg.jinjaValue()
         ])
 
         #expect(got == exp)