Tests: reuse a couple of tokenizers (#221)

* Tests: reuse a couple of tokenizers

* I hate formatters

Author: pcuenca

Tests/TokenizersTests/ChatTemplateTests.swift

@@ -14,8 +14,24 @@ class ChatTemplateTests: XCTestCase {
         "content": "Describe the Swift programming language.",
     ]]
 
+    static let phiTokenizerTask = Task {
+        try await AutoTokenizer.from(pretrained: "microsoft/Phi-3-mini-128k-instruct")
+    }
+
+    static func sharedPhiTokenizer() async throws -> Tokenizer {
+        try await phiTokenizerTask.value
+    }
+
+    static let tokenizerWithTemplateArrayTask = Task {
+        try await AutoTokenizer.from(pretrained: "mlx-community/Mistral-7B-Instruct-v0.3-4bit")
+    }
+
+    static func sharedTokenizerWithTemplateArray() async throws -> Tokenizer {
+        try await tokenizerWithTemplateArrayTask.value
+    }
+
     func testTemplateFromConfig() async throws {
-        let tokenizer = try await AutoTokenizer.from(pretrained: "microsoft/Phi-3-mini-128k-instruct")
+        let tokenizer = try await Self.sharedPhiTokenizer()
         let encoded = try tokenizer.applyChatTemplate(messages: messages)
         let encodedTarget = [32010, 4002, 29581, 278, 14156, 8720, 4086, 29889, 32007, 32001]
         let decoded = tokenizer.decode(tokens: encoded)
@@ -36,7 +52,7 @@ class ChatTemplateTests: XCTestCase {
     }
 
     func testDefaultTemplateFromArrayInConfig() async throws {
-        let tokenizer = try await AutoTokenizer.from(pretrained: "mlx-community/Mistral-7B-Instruct-v0.3-4bit")
+        let tokenizer = try await Self.sharedTokenizerWithTemplateArray()
         let encoded = try tokenizer.applyChatTemplate(messages: messages)
         let encodedTarget = [1, 29473, 3, 28752, 1040, 4672, 2563, 17060, 4610, 29491, 29473, 4]
         let decoded = tokenizer.decode(tokens: encoded)
@@ -46,7 +62,7 @@ class ChatTemplateTests: XCTestCase {
     }
 
     func testTemplateFromArgumentWithEnum() async throws {
-        let tokenizer = try await AutoTokenizer.from(pretrained: "microsoft/Phi-3-mini-128k-instruct")
+        let tokenizer = try await Self.sharedPhiTokenizer()
         // Purposely not using the correct template for this model to verify that the template from the config is not being used
         let mistral7BDefaultTemplate = "{{bos_token}}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ ' [INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + ' ' + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"
         let encoded = try tokenizer.applyChatTemplate(messages: messages, chatTemplate: .literal(mistral7BDefaultTemplate))
@@ -58,7 +74,7 @@ class ChatTemplateTests: XCTestCase {
     }
 
     func testTemplateFromArgumentWithString() async throws {
-        let tokenizer = try await AutoTokenizer.from(pretrained: "microsoft/Phi-3-mini-128k-instruct")
+        let tokenizer = try await Self.sharedPhiTokenizer()
         // Purposely not using the correct template for this model to verify that the template from the config is not being used
         let mistral7BDefaultTemplate = "{{bos_token}}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ ' [INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + ' ' + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"
         let encoded = try tokenizer.applyChatTemplate(messages: messages, chatTemplate: mistral7BDefaultTemplate)
@@ -70,7 +86,7 @@ class ChatTemplateTests: XCTestCase {
     }
 
     func testNamedTemplateFromArgument() async throws {
-        let tokenizer = try await AutoTokenizer.from(pretrained: "mlx-community/Mistral-7B-Instruct-v0.3-4bit")
+        let tokenizer = try await Self.sharedTokenizerWithTemplateArray()
         // Normally it is not necessary to specify the name `default`, but I'm not aware of models with lists of templates in the config that are not `default` or `tool_use`
         let encoded = try tokenizer.applyChatTemplate(messages: messages, chatTemplate: .name("default"))
         let encodedTarget = [1, 29473, 3, 28752, 1040, 4672, 2563, 17060, 4610, 29491, 29473, 4]