Merge pull request #21 from CodeWithKyrian/add-chat-mode-detection

Add chat input detection for text generation, and refactor streamer API.

Author: CodeWithKyrian

examples/pipelines/summarization.php

@@ -14,7 +14,7 @@
 
 $summarizer = pipeline('summarization', 'Xenova/distilbart-cnn-6-6');
 
-$streamer = StdOutStreamer::make($summarizer->tokenizer);
+$streamer = StdOutStreamer::make();
 
 $article = 'The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, ' .
     'and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. ' .
@@ -29,6 +29,6 @@
 //$article = "I called my friend to see if he wanted to go to the movies. However, he was busy and said he would " .
 //    "call me back. I didn't hear back from him, so I called him again. He didn't answer, so I decided to go to the movies by myself.";
 
-$summary = $summarizer($article, streamer: $streamer,  maxNewTokens: 512, temperature: 0.7);
+$summary = $summarizer($article, streamer: $streamer, maxNewTokens: 512, temperature: 0.7);
 
 dd("Done", timeUsage(), memoryUsage());

examples/pipelines/text-generation.php

@@ -14,19 +14,20 @@
 //$generator = pipeline('text-generation', 'Xenova/gpt2');
 $generator = pipeline('text-generation', 'Xenova/Qwen1.5-0.5B-Chat');
 
-$streamer = StdOutStreamer::make($generator->tokenizer);
+$streamer = StdOutStreamer::make();
 
 $messages = [
     ['role' => 'system', 'content' => 'You are a helpful assistant.'],
-    ['role' => 'user', 'content' => 'Who are you'],
+    ['role' => 'user', 'content' => 'What is the product of 5 and 4'],
 ];
 
 $input = $generator->tokenizer->applyChatTemplate($messages, addGenerationPrompt: true, tokenize: false);
 
-$output = $generator($messages,
+$output = $generator($input,
     streamer: $streamer,
     maxNewTokens: 128,
     doSample: true,
+    returnFullText: false,
 //    temperature: 0.7,
 //    repetitionPenalty: 1.3,
 //    earlyStopping: true
@@ -42,4 +43,4 @@
 //    doSample: true
 //);
 
-dd("done", timeUsage(), memoryUsage());
+dd($output[0]['generated_text'], timeUsage(), memoryUsage());

examples/pipelines/text2text-generation.php

@@ -12,7 +12,7 @@
 $generator = pipeline('text2text-generation', 'Xenova/LaMini-Flan-T5-783M');
 //$generator = pipeline('text2text-generation', 'Xenova/flan-t5-small', quantized: true);
 
-$streamer = StdOutStreamer::make($generator->tokenizer);
+$streamer = StdOutStreamer::make();
 
 //$query = 'Please let me know your thoughts on the given place and why you think it deserves to be visited: \n" Barcelona, Spain"';
 $query = 'How many continents are in the world? List them out';

examples/pipelines/translation.php

@@ -13,10 +13,10 @@
 
 $translator = pipeline('translation', 'Xenova/m2m100_418M');
 
-$streamer = StdOutStreamer::make($translator->tokenizer);
+$streamer = StdOutStreamer::make();
 
-$output = $translator('生活就像一盒巧克力。', streamer: $streamer, tgtLang: 'en');
+//$output = $translator('生活就像一盒巧克力。', streamer: $streamer, tgtLang: 'en');
 //$output = $translator('जीवन एक चॉकलेट बॉक्स की तरह है।', streamer: $streamer, tgtLang: 'fr');
-//$output = $translator('संयुक्त राष्ट्र के प्रमुख का कहना है कि सीरिया में कोई सैन्य समाधान नहीं है', streamer: $streamer, tgtLang: 'fr', maxNewTokens: 256);
+$output = $translator('संयुक्त राष्ट्र के प्रमुख का कहना है कि सीरिया में कोई सैन्य समाधान नहीं है', streamer: $streamer, tgtLang: 'fr', maxNewTokens: 256);
 
 dd("done", timeUsage(), memoryUsage());

src/Generation/Streamers/StdOutStreamer.php

@@ -5,18 +5,16 @@
 
 namespace Codewithkyrian\Transformers\Generation\Streamers;
 
-use Codewithkyrian\Transformers\PretrainedTokenizers\PretrainedTokenizer;
-
 class StdOutStreamer extends TextStreamer
 {
-    public function __construct(PretrainedTokenizer $tokenizer, StreamMode $streamMode = StreamMode::PARTIAL)
+    public function __construct(StreamMode $streamMode = StreamMode::PARTIAL)
     {
-        parent::__construct($tokenizer);
+        parent::__construct();
+
         $this->setStreamMode($streamMode);
-        if($streamMode === StreamMode::FULL) {
+        if ($streamMode === StreamMode::FULL) {
             $this->onStreamEnd(fn(string $full) => $this->echoToConsole($full, true));
-        }
-        else {
+        } else {
             $this->onStream(fn(string $partial) => $this->echoToConsole($partial));
             $this->onStreamEnd(fn(string $full) => $this->echoToConsole('', true));
         }

src/Generation/Streamers/Streamer.php

@@ -5,11 +5,15 @@
 
 namespace Codewithkyrian\Transformers\Generation\Streamers;
 
+use Codewithkyrian\Transformers\PretrainedTokenizers\PretrainedTokenizer;
+
 /**
  * Base streamer from which all streamers inherit.
  */
 abstract class Streamer
 {
+    abstract public function init(PretrainedTokenizer $tokenizer, array $inputTokens, bool $excludeInput = false): void;
+
     abstract public function put(mixed $value): void;
 
     abstract public function end(): void;

src/Generation/Streamers/TextStreamer.php

@@ -88,6 +88,9 @@
  */
 class TextStreamer extends Streamer
 {
+    protected PretrainedTokenizer $tokenizer;
+    protected array $inputTokens = [];
+    protected bool $excludeInput = true;
     protected string $printedText = '';
     protected mixed $onStreamCallback = null;
     protected mixed $onStreamEndCallback = null;
@@ -96,13 +99,28 @@ class TextStreamer extends Streamer
     protected int $lastDecodedCheckpointForToken = 0;
     protected int $lastDecodedCheckpointForText = 0;
 
-    public function __construct(protected PretrainedTokenizer $tokenizer)
+    public function __construct()
     {
     }
 
-    public static function make(PretrainedTokenizer $tokenizer): self
+    public static function make(): self
     {
-        return new static($tokenizer);
+        return new static();
+    }
+
+    public function init(PretrainedTokenizer $tokenizer, array $inputTokens, bool $excludeInput = false): void
+    {
+        $this->tokenizer = $tokenizer;
+        $this->inputTokens = $inputTokens;
+        $this->excludeInput = $excludeInput;
+
+        if ($this->excludeInput) {
+            $this->printedText = $this->tokenizer->decode($this->inputTokens, skipSpecialTokens: true);
+            $this->printedLength = mb_strlen($this->printedText);
+
+            $this->lastDecodedCheckpointForToken = count($this->inputTokens);
+            $this->lastDecodedCheckpointForText = mb_strlen($this->printedText);
+        }
     }
 
     public function onStream(callable $callback): self

src/Pipelines/Text2TextGenerationPipeline.php

@@ -5,8 +5,8 @@
 
 namespace Codewithkyrian\Transformers\Pipelines;
 
+use Codewithkyrian\Transformers\Generation\Streamers\Streamer;
 use Codewithkyrian\Transformers\Utils\GenerationConfig;
-use Codewithkyrian\Transformers\Utils\Tensor;
 
 /**
  * A pipeline for generating text using a model that performs text-to-text generation tasks.
@@ -31,6 +31,7 @@ public function __invoke(array|string $inputs, ...$args): array
         $streamer = null;
 
         if (array_key_exists('streamer', $args)) {
+            /** @var Streamer $streamer */
             $streamer = $args['streamer'];
             unset($args['streamer']);
         }
@@ -77,6 +78,9 @@ public function __invoke(array|string $inputs, ...$args): array
             : $tokenizer->__invoke($inputs, padding: true, truncation: true)['input_ids'];
 
 
+        // Streamer can only handle one input at a time for now, so we only pass the first input
+        $streamer->init($this->tokenizer, $inputIds[0]->toArray());
+
         // Generate output token ids
         $outputTokenIds = $this->model->generate($inputIds, generationConfig: $generateKwargs, streamer: $streamer);
 

src/Pipelines/TextGenerationPipeline.php

@@ -5,7 +5,9 @@
 
 namespace Codewithkyrian\Transformers\Pipelines;
 
+use Codewithkyrian\Transformers\Generation\Streamers\Streamer;
 use Codewithkyrian\Transformers\Utils\GenerationConfig;
+use function Codewithkyrian\Transformers\Utils\array_every;
 use function Codewithkyrian\Transformers\Utils\camelCaseToSnakeCase;
 
 /**
@@ -53,12 +55,18 @@ class TextGenerationPipeline extends Pipeline
     public function __invoke(array|string $inputs, ...$args): array
     {
         $streamer = null;
-
         if (array_key_exists('streamer', $args)) {
+            /** @var Streamer $streamer */
             $streamer = $args['streamer'];
             unset($args['streamer']);
         }
 
+        $returnFullText = true; // By default, return full text
+        if (array_key_exists('returnFullText', $args)) {
+            $returnFullText = $args['returnFullText'];
+            unset($args['returnFullText']);
+        }
+
         // Convert the rest of the arguments key names from camelCase to snake_case
         $snakeCasedArgs = [];
         foreach ($args as $key => $value) {
@@ -67,45 +75,79 @@ public function __invoke(array|string $inputs, ...$args): array
 
         $generationConfig = new GenerationConfig($snakeCasedArgs);
 
-        $isChatMode = $this->isChatMode($inputs);
-
-        if ($isChatMode) {
-            $inputs = $this->tokenizer->applyChatTemplate($inputs, addGenerationPrompt: true, tokenize: false);
-        }
-
-        $isBatched = is_array($inputs);
-
-        if (!$isBatched) {
-            $inputs = [$inputs];
+        $isBatched = false;
+        $isChatInput = false;
+
+        // Normalize inputs
+        $texts = [];
+
+        if (is_string($inputs)) {
+            $texts = $inputs = [$inputs];
+        } elseif (is_array($inputs) && array_every($inputs, fn($x) => is_string($x))) {
+            $isBatched = true;
+            $texts = $inputs;
+        } else {
+            if ($this->isChat($inputs)) {
+                $inputs = [$inputs];
+            } elseif (is_array($inputs) && array_every($inputs, [$this, 'isChat'])) {
+                $isBatched = true;
+            } else {
+                throw new \Exception('Input must be a string, an array of strings, a Chat, or an array of Chats');
+            }
+            $isChatInput = true;
+
+            // If the input is a chat, apply the chat template
+            $texts = array_map(fn($x) => $this->tokenizer->applyChatTemplate($x, addGenerationPrompt: true, tokenize: false), $inputs);
         }
 
         // By default, do not add special tokens
-        $addSpecialTokens = $this->model->config['add_special_tokens'] ?? false;
+        $addSpecialTokens = $generationConfig['add_special_tokens'] ?? false;
+
+        $returnFullText = $isChatInput ? false : $returnFullText;
 
         $this->tokenizer->paddingSide = 'left';
         ['input_ids' => $inputIds, 'attention_mask' => $attentionMask] = $this->tokenizer->tokenize(
-            $inputs,
+            $texts,
             padding: true,
             addSpecialTokens: $addSpecialTokens,
             truncation: true
         );
 
-        $outputTokenIds = $this->model->generate(
-            $inputIds,
+        // Streamer can only handle one input at a time for now, so we only pass the first input
+        $streamer->init($this->tokenizer, $inputIds[0]->toArray(), true);
+
+        $outputTokenIds = $this->model->generate($inputIds,
             generationConfig: $generationConfig,
             inputsAttentionMask: $attentionMask,
             streamer: $streamer
         );
 
         $decoded = $this->tokenizer->batchDecode($outputTokenIds, skipSpecialTokens: true);
 
+        $promptLengths = null;
+        if (!$returnFullText && $inputIds->shape()[count($inputIds->shape()) - 1] > 0) {
+            $promptLengths = array_map(fn($x) => mb_strlen($x), $this->tokenizer->batchDecode($inputIds->toArray(), skipSpecialTokens: true));
+        }
 
         $toReturn = array_fill(0, count($inputs), []);
 
         for ($i = 0; $i < count($decoded); ++$i) {
             $textIndex = floor($i / count($outputTokenIds) * count($inputs));
+
+            if ($promptLengths !== null) {
+                // Trim the decoded text to only include the generated part
+                $decoded[$i] = substr($decoded[$i], $promptLengths[$textIndex]);
+
+                // Remove the leading space
+                $decoded[$i] = ltrim($decoded[$i]);
+            }
+
             $toReturn[$textIndex][] = [
-                'generated_text' => $decoded[$i]
+                'generated_text' => $isChatInput
+                    ? array_merge($inputs[$textIndex], [
+                        ['role' => 'assistant', 'content' => $decoded[$i]]
+                    ])
+                    : $decoded[$i],
             ];
         }
 
@@ -114,9 +156,9 @@ public function __invoke(array|string $inputs, ...$args): array
     }
 
     // Detect chat mode
-    protected function isChatMode(string|array $texts): bool
+    function isChat($x): bool
     {
-        return is_array($texts) && isset($texts[0]) && is_array($texts[0]) && !array_is_list($texts[0]);
-
+        return is_array($x) && array_every($x, fn($item) => isset($item['role']) && isset($item['content']));
     }
+
 }