Optimize Tensor softmax for 1 dimensional tensors

Author: CodeWithKyrian

examples/pipelines/text-generation.php

@@ -13,35 +13,35 @@
 //
 //$generator = pipeline('text-generation', 'Xenova/gpt2');
 //$generator = pipeline('text-generation', 'Xenova/Qwen1.5-0.5B-Chat');
-//$generator = pipeline('text-generation', 'Xenova/TinyLlama-1.1B-Chat-v1.0');
-//
-//$streamer = StdOutStreamer::make();
-//
-//$messages = [
-//    ['role' => 'system', 'content' => 'You are a helpful assistant.'],
-//    ['role' => 'user', 'content' => 'What is diffusion in chemistry?'],
-//];
-//
-//$input = $generator->tokenizer->applyChatTemplate($messages, addGenerationPrompt: true, tokenize: false);
-//
-//$output = $generator($input,
-//    streamer: $streamer,
-//    maxNewTokens: 128,
-//    doSample: true,
-//    returnFullText: false,
-////    temperature: 0.7,
-////    repetitionPenalty: 1.3,
-////    earlyStopping: true
-//);
+$generator = pipeline('text-generation', 'Xenova/TinyLlama-1.1B-Chat-v1.0');
 
-$generator = pipeline('text-generation', 'Xenova/codegen-350M-mono');
-$streamer = StdOutStreamer::make($generator->tokenizer);
+$streamer = StdOutStreamer::make();
 
-$output = $generator(
-    'def fib(n):',
+$messages = [
+    ['role' => 'system', 'content' => 'You are a helpful assistant.'],
+    ['role' => 'user', 'content' => 'What is diffusion?'],
+];
+
+$input = $generator->tokenizer->applyChatTemplate($messages, addGenerationPrompt: true, tokenize: false);
+
+$output = $generator($input,
     streamer: $streamer,
-    maxNewTokens: 100,
-    doSample: true
+    maxNewTokens: 256,
+    doSample: true,
+    returnFullText: false,
+//    temperature: 0.7,
+//    repetitionPenalty: 1.3,
+//    earlyStopping: true
 );
 
+//$generator = pipeline('text-generation', 'Xenova/codegen-350M-mono');
+//$streamer = StdOutStreamer::make($generator->tokenizer);
+//
+//$output = $generator(
+//    'def fib(n):',
+//    streamer: $streamer,
+//    maxNewTokens: 100,
+//    doSample: true
+//);
+
 dd($output[0]['generated_text'], timeUsage(), memoryUsage());

src/Generation/Samplers/Sampler.php

@@ -47,7 +47,7 @@ abstract public function sample(Tensor $logits, int $index);
     public function getLogits(Tensor $logits, int $index): array
     {
         $vocabSize = $logits->shape()[count($logits->shape()) - 1];
-        $logs = $logits->buffer()->toArray();
+        $logs = $logits->toBufferArray();
 
         if ($index === -1) {
             $logs = array_slice($logs, -$vocabSize);
@@ -76,7 +76,6 @@ public function randomSelect(array $probabilities): int
 
         // Generate a random number between 0 and the sum of probabilities
         $r = mt_rand() / mt_getrandmax() * $sumProbabilities;
-
         foreach ($probabilities as $i => $probability) {
             $r -= $probability;
 

src/Generation/Streamers/TextStreamer.php

@@ -89,7 +89,6 @@
 class TextStreamer extends Streamer
 {
     protected PretrainedTokenizer $tokenizer;
-    protected array $inputTokens = [];
     protected bool $excludeInput = true;
     protected string $printedText = '';
     protected mixed $onStreamCallback = null;
@@ -111,14 +110,14 @@ public static function make(): self
     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->printedText = $this->tokenizer->decode($inputTokens, skipSpecialTokens: true);
             $this->printedLength = mb_strlen($this->printedText);
 
-            $this->lastDecodedCheckpointForToken = count($this->inputTokens);
+
+            $this->lastDecodedCheckpointForToken = count($inputTokens) - 1;
             $this->lastDecodedCheckpointForText = mb_strlen($this->printedText);
         }
     }
@@ -154,7 +153,6 @@ public function put(mixed $value): void
         // Check for punctuation marks indicating the end of a word or sentence
         $punctuationMarks = ['.', ',', '!', '?', ';', ':'];
 
-
         $this->printedText = mb_substr($this->printedText, 0, $this->lastDecodedCheckpointForText)
             . ($this->lastDecodedCheckpointForToken == 0 ? '' : ' ')
             . $decodedText;

src/Utils/Tensor.php

@@ -974,7 +974,7 @@ public function permute(...$axes): static
     public function softmax(): array|static
     {
         return match ($this->ndim()) {
-            1 => Math::softmax($this->toArray()),
+            1 => $this->unsqueeze(0)->softmax2D(),
             2 => $this->softmax2D(),
             default => throw new InvalidArgumentException("Softmax is only supported for 1D and 2D tensors.")
         };