Merge pull request #1 from CodeWithKyrian/apply-chat-template

Add ability to use chat templates in Text Generation

Author: CodeWithKyrian

composer.json

@@ -18,7 +18,8 @@
     "ankane/onnxruntime": "^0.2.0",
     "ext-gd": "*",
     "rindow/rindow-math-matrix": "^1.2",
-    "spatie/fork": "^1.2"
+    "spatie/fork": "^1.2",
+    "codewithkyrian/jinja-php": "^1.0"
   },
   "require-dev": {
     "pestphp/pest": "^2.31",

examples/tokenizers/apply-chat-template.php

@@ -0,0 +1,20 @@
+<?php
+
+declare(strict_types=1);
+
+use Codewithkyrian\Transformers\PretrainedTokenizers\AutoTokenizer;
+
+require_once './bootstrap.php';
+
+$tokenizer = AutoTokenizer::fromPretrained('mistralai/Mistral-7B-Instruct-v0.1');
+//$tokenizer = AutoTokenizer::fromPretrained('facebook/blenderbot-400M-distill');
+$messages = [
+    ['role' => 'user', 'content' => 'Hello!'],
+    ['role' => 'assistant', 'content' => 'Hi! How are you?'],
+    ['role' => 'user', 'content' => 'I am doing great.'],
+    ['role' => 'assistant', 'content' => 'That is great to hear.'],
+];
+
+$text = $tokenizer->applyChatTemplate($messages, addGenerationPrompt: true, tokenize: false);
+
+dd($text);

src/Exceptions/TemplateParseException.php

@@ -0,0 +1,14 @@
+<?php
+
+declare(strict_types=1);
+
+
+namespace Codewithkyrian\Transformers\Exceptions;
+
+class TemplateParseException extends \Exception implements TransformersException
+{
+    public static function undefinedVariable($variableName): TemplateParseException
+    {
+        return new self("Undefined variable:  $variableName");
+    }
+}

src/Pipelines/TextGenerationPipeline.php

@@ -66,7 +66,14 @@ public function __invoke(array|string $texts, ...$args): array
 
         $generationConfig = new GenerationConfig($snakeCasedArgs);
 
+        $isChatMode = $this->isChatMode($texts);
+
+        if ($isChatMode) {
+            $texts = $this->tokenizer->applyChatTemplate($texts, addGenerationPrompt: true, tokenize: false);
+        }
+
         $isBatched = is_array($texts);
+
         if (!$isBatched) {
             $texts = [$texts];
         }
@@ -86,6 +93,7 @@ public function __invoke(array|string $texts, ...$args): array
 
         $decoded = $this->tokenizer->batchDecode($outputTokenIds, skipSpecialTokens: true);
 
+
         $toReturn = array_fill(0, count($texts), []);
 
         for ($i = 0; $i < count($decoded); ++$i) {
@@ -103,4 +111,11 @@ protected function camelCaseToSnakeCase(string $input): string
     {
         return strtolower(preg_replace('/(?<!^)[A-Z]/', '_$0', $input));
     }
+
+    // Detect chat mode
+    protected function isChatMode(string|array $texts): bool
+    {
+        return is_array($texts) && isset($texts[0]) && is_array($texts[0]) && !array_is_list($texts[0]);
+
+    }
 }

src/PretrainedTokenizers/PretrainedTokenizer.php

@@ -5,6 +5,7 @@
 
 namespace Codewithkyrian\Transformers\PretrainedTokenizers;
 
+use Codewithkyrian\Jinja\Template;
 use Codewithkyrian\Transformers\Decoders\Decoder;
 use Codewithkyrian\Transformers\Normalizers\Normalizer;
 use Codewithkyrian\Transformers\PostProcessors\PostProcessedOutput;
@@ -18,7 +19,7 @@ class PretrainedTokenizer
 {
     protected bool $returnTokenTypeIds = false;
 
-    protected bool $warnedAboutChatTemplate;
+    protected bool $warnedAboutChatTemplate = false;
 
     protected string $defaultChatTemplate = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}";
 
@@ -68,7 +69,7 @@ class PretrainedTokenizer
     protected bool $legacy;
 
     protected mixed $chatTemplate;
-    protected \SplObjectStorage $compiledTemplateCache;
+    protected array $compiledTemplateCache = [];
 
     /**
      * @param array $tokenizerJSON The JSON of the tokenizer.
@@ -145,8 +146,6 @@ public function __construct(protected array $tokenizerJSON, protected array $tok
         $this->legacy = false;
 
         $this->chatTemplate = $tokenizerConfig['chat_template'] ?? null;
-        $this->compiledTemplateCache = new \SplObjectStorage();
-
     }
 
     /**
@@ -335,7 +334,7 @@ function ($key) {
             if (
                 array_reduce($encodedTokens, function ($carry, $x) use ($encodedTokens) {
                     foreach ($x as $key => $value) {
-                        if (count($value) !== count($encodedTokens[0][$key] ?? [])) {
+                        if (count($value ?? []) !== count($encodedTokens[0][$key] ?? [])) {
                             return true;
                         }
                     }
@@ -610,10 +609,10 @@ function truncateHelper(array &$item, int $length): void
 
     protected function getDefaultChatTemplate(): string
     {
-        if (!$this->warnedAboutChatTemplate) {
-            trigger_error("The default chat template is deprecated and will be removed in a future version. Please use the `chat_template` option instead.", E_USER_WARNING);
-            $this->warnedAboutChatTemplate = true;
-        }
+//        if (!$this->warnedAboutChatTemplate) {
+//            trigger_error("The default chat template is deprecated and will be removed in a future version. Please use the `chat_template` option instead.", E_USER_WARNING);
+//            $this->warnedAboutChatTemplate = true;
+//        }
 
         return $this->defaultChatTemplate;
     }
@@ -670,7 +669,7 @@ public function applyChatTemplate(
 
         if ($compiledTemplate === null) {
             // TODO: Use Jinja to compile the template
-            $compiledTemplate = null;
+            $compiledTemplate =  new Template($chatTemplate);
             $this->compiledTemplateCache[$chatTemplate] = $compiledTemplate;
         }
 
@@ -695,10 +694,10 @@ public function applyChatTemplate(
                 addSpecialTokens: false,
                 truncation: $truncation,
                 maxLength: $maxLength
-            )['input_ids'];
+            )['input_ids']->toArray();
         }
 
-        return $rendered;
+        return stripcslashes($rendered);
     }
 
 }

src/Utils/Helpers.php

@@ -58,7 +58,6 @@ function array_every(array $array, callable $callback): bool
     return true;
 }
 
-
 function joinPaths(string ...$args): string
 {
     $paths = [];