Add Qwen2 model support

docs/getting-started.md

@@ -12,6 +12,7 @@ Before installing TransformersPHP, ensure your system meets the following requir
 - Composer
 - PHP FFI extension
 - JIT compilation (optional)
+- Increased memory limit (for advanced tasks like text generation)
 
 ## Installation
 

examples/pipelines/image-to-text.php

@@ -9,8 +9,8 @@
 require_once './bootstrap.php';
 
 ini_set('memory_limit', -1);
-$captioner = pipeline('image-to-text', 'Xenova/vit-gpt2-image-captioning');
-//$captioner = pipeline('image-to-text', 'Xenova/trocr-small-handwritten');
+//$captioner = pipeline('image-to-text', 'Xenova/vit-gpt2-image-captioning');
+$captioner = pipeline('image-to-text', 'Xenova/trocr-small-handwritten');
 
 //$streamer = StdOutStreamer::make($captioner->tokenizer);
 

examples/pipelines/text-generation.php

@@ -11,23 +11,25 @@
 
 ini_set('memory_limit', -1);
 //
-$generator = pipeline('text-generation', 'Xenova/gpt2');
+//$generator = pipeline('text-generation', 'Xenova/gpt2');
+$generator = pipeline('text-generation', 'Xenova/Qwen1.5-0.5B-Chat');
 
 $streamer = StdOutStreamer::make($generator->tokenizer);
 
 $messages = [
-    ['role' => 'user', 'content' => 'Hello!'],
-    ['role' => 'assistant', 'content' => 'Hi! How are you?'],
-    ['role' => 'user', 'content' => 'I am doing great. What about you?'],
+    ['role' => 'system', 'content' => 'You are a helpful assistant.'],
+    ['role' => 'user', 'content' => 'Who are you'],
 ];
 
-$output = $generator("I love going to school but I don't",
+$input = $generator->tokenizer->applyChatTemplate($messages, addGenerationPrompt: true, tokenize: false);
+
+$output = $generator($messages,
     streamer: $streamer,
     maxNewTokens: 128,
     doSample: true,
-    temperature: 0.7,
-    repetitionPenalty: 1.3,
-    earlyStopping: true
+//    temperature: 0.7,
+//    repetitionPenalty: 1.3,
+//    earlyStopping: true
 );
 
 //$generator = pipeline('text-generation', 'Xenova/codegen-350M-mono');

src/Models/Auto/AutoModel.php

@@ -38,6 +38,7 @@ class AutoModel extends PretrainedMixin
         "gptj" => \Codewithkyrian\Transformers\Models\Pretrained\GPTJModel::class,
         "gpt_bigcode" => \Codewithkyrian\Transformers\Models\Pretrained\GPTBigCodeModel::class,
         "codegen" => \Codewithkyrian\Transformers\Models\Pretrained\CodeGenModel::class,
+        "qwen2" => \Codewithkyrian\Transformers\Models\Pretrained\Qwen2Model::class,
     ];
 
     const MODEL_CLASS_MAPPINGS = [

src/Models/Auto/AutoModelForCausalLM.php

@@ -13,7 +13,7 @@ class AutoModelForCausalLM extends PretrainedMixin
         'gpt_bigcode' => \Codewithkyrian\Transformers\Models\Pretrained\GPTBigCodeForCausalLM::class,
         'codegen' => \Codewithkyrian\Transformers\Models\Pretrained\CodeGenForCausalLM::class,
         'trocr' => \Codewithkyrian\Transformers\Models\Pretrained\TrOCRForCausalLM::class,
-
+        'qwen2' => \Codewithkyrian\Transformers\Models\Pretrained\Qwen2ForCausalLM::class
     ];
 
     const MODEL_CLASS_MAPPINGS = [

src/Models/ModelArchitecture.php

@@ -9,6 +9,7 @@
 use Codewithkyrian\Transformers\Models\Pretrained\PretrainedModel;
 use Codewithkyrian\Transformers\Utils\GenerationConfig;
 use Codewithkyrian\Transformers\Utils\Tensor;
+use Interop\Polite\Math\Matrix\NDArray;
 
 enum ModelArchitecture: string
 {
@@ -34,7 +35,7 @@ public function runBeam(PretrainedModel $model, array &$beam): array
     {
         return match ($this) {
             self::DecoderOnly => $this->decoderRunBeam($model, $beam),
-            self::Seq2SeqLM, self::Vision2Seq  => $this->seq2seqRunBeam($model, $beam),
+            self::Seq2SeqLM, self::Vision2Seq => $this->seq2seqRunBeam($model, $beam),
             default => throw new \Error('This model type does not support beam search'),
         };
     }
@@ -114,10 +115,11 @@ protected function decoderRunBeam(PretrainedModel $model, array &$beam): array
         // 1. Prepare
         $modelInputs = [
             'input_ids' => $beam['model_input_ids'],
-            'attention_mask' => new Tensor($attnMaskData, shape: [1, $attnMaskLength]),
+            'attention_mask' => new Tensor($attnMaskData, NDArray::int64, [1, $attnMaskLength]),
             'past_key_values' => $beam['prev_model_outputs']['past_key_values'] ?? null,
         ];
 
+
         // 2. Run
         $output = $model->forward($modelInputs);
 
@@ -155,7 +157,7 @@ protected function decoderStartBeams(
             $attnMask = null;
             if ($inputsAttentionMask !== null) {
                 $attnMask = $inputsAttentionMask[$beamId];
-                $attnMask->reshape([1, ...$attnMask->shape()]);
+                $attnMask = $attnMask->reshape([1, ...$attnMask->shape()]);
             } else {
                 $attnMask = $model->prepareAttentionMask($tokens);
             }
@@ -189,8 +191,7 @@ protected function decoderStartBeams(
     protected function decoderUpdatebeam(array &$beam, int $newTokenId): void
     {
         $beam['output_token_ids'][] = $newTokenId;
-
-        $beam['model_input_ids'] = new Tensor([$newTokenId], shape: [1, 1]);
+        $beam['model_input_ids'] = new Tensor([$newTokenId], NDArray::int64, [1, 1]);
     }
 
     /**
@@ -221,6 +222,14 @@ protected function decoderForward(PretrainedModel $model, array $modelInputs): a
         $model->preparePositionIds($inputNames, $decoderFeeds, $useCacheBranch);
         $model->addPastKeyValues($decoderFeeds, $pastKeyValues);
 
+        // The initial past key values should have a shape of 0 in one of the dimensions, which
+        // is the sequence length. However, I haven't found a way to pass a tensor with a shape of 0
+        // to the model, so I'm using a sequence length of 1 instead for the first step, and then
+        // offsetting the sequence length by 1 for the subsequent steps. This is a workaround for now.
+        $prevSequenceLength = $decoderFeeds['past_key_values.0.key']->shape()[2];
+        $attnMaskLength = $prevSequenceLength == 1 ? 1 : $prevSequenceLength + 1;
+        $decoderFeeds['attention_mask'] = Tensor::ones([1, $attnMaskLength], dtype: NDArray::int64);
+
         $decoderResults = $model->runSession($model->session, $decoderFeeds);
 
         $logits = $decoderResults['logits'];

src/Models/Pretrained/PretrainedModel.php

@@ -399,7 +399,7 @@ public function preparePositionIds(array $inputNames, array &$feeds, bool $useCa
 
         if ($useCacheBranch) {
             // TODO: Fix this
-            $feeds['position_ids'] = $feeds['position_ids']->slice(null, -1)->unsqueeze(-1);
+            $feeds['position_ids'] = $feeds['position_ids']->slice(null, -1);
         }
     }
 

src/Models/Pretrained/Qwen2ForCausalLM.php

@@ -0,0 +1,11 @@
+<?php
+
+declare(strict_types=1);
+
+
+namespace Codewithkyrian\Transformers\Models\Pretrained;
+
+class Qwen2ForCausalLM extends Qwen2PreTrainedModel
+{
+
+}

src/Models/Pretrained/Qwen2Model.php

@@ -0,0 +1,11 @@
+<?php
+
+declare(strict_types=1);
+
+
+namespace Codewithkyrian\Transformers\Models\Pretrained;
+
+class Qwen2Model extends Qwen2PreTrainedModel
+{
+
+}

src/Models/Pretrained/Qwen2PreTrainedModel.php

@@ -0,0 +1,39 @@
+<?php
+
+declare(strict_types=1);
+
+
+namespace Codewithkyrian\Transformers\Models\Pretrained;
+
+use Codewithkyrian\Transformers\Models\ModelArchitecture;
+use Codewithkyrian\Transformers\Utils\AutoConfig;
+use Codewithkyrian\Transformers\Utils\GenerationConfig;
+use OnnxRuntime\InferenceSession;
+
+/**
+ * The bare Qwen2 Model outputting raw hidden-states without any specific head on top.
+ */
+class Qwen2PreTrainedModel extends PreTrainedModel
+{
+    protected int $numHeads;
+    protected int $numLayers;
+    protected int $dimKv;
+
+    public function __construct(
+        AutoConfig               $config,
+        InferenceSession         $session,
+        public ModelArchitecture $modelArchitecture,
+        public GenerationConfig  $generationConfig
+    )
+    {
+        parent::__construct($config, $session, $modelArchitecture);
+
+        // config doesn't contain pad_token_id, so we assume it is the eos_token_id
+        $this->config['pad_token_id'] = $this->config['eos_token_id'];
+        $this->config->padTokenId = $this->config['eos_token_id'];
+
+        $this->numHeads = $this->config['num_key_value_heads'] ??  $this->config['num_attention_heads'];
+        $this->numLayers = $this->config['num_hidden_layers'];
+        $this->dimKv = $this->config['hidden_size'] / $this->config['num_attention_heads'];
+    }
+}

src/Pipelines/TextGenerationPipeline.php

@@ -90,7 +90,12 @@ public function __invoke(array|string $inputs, ...$args): array
             truncation: true
         );
 
-        $outputTokenIds = $this->model->generate($inputIds, generationConfig: $generationConfig, streamer: $streamer);
+        $outputTokenIds = $this->model->generate(
+            $inputIds,
+            generationConfig: $generationConfig,
+            inputsAttentionMask: $attentionMask,
+            streamer: $streamer
+        );
 
         $decoded = $this->tokenizer->batchDecode($outputTokenIds, skipSpecialTokens: true);
 

src/PretrainedTokenizers/AutoTokenizer.php

@@ -55,6 +55,7 @@ class AutoTokenizer
         'SpeechT5Tokenizer' => 'SpeechT5Tokenizer',
         'NougatTokenizer' => 'NougatTokenizer',
         'VitsTokenizer' => 'VitsTokenizer',
+        'Qwen2Tokenizer' => 'Qwen2Tokenizer',
         // Base case:
         'PreTrainedTokenizer' => 'PreTrainedTokenizer',
     ];

src/PretrainedTokenizers/Qwen2Tokenizer.php

@@ -0,0 +1,11 @@
+<?php
+
+declare(strict_types=1);
+
+
+namespace Codewithkyrian\Transformers\PretrainedTokenizers;
+
+class Qwen2Tokenizer extends PretrainedTokenizer
+{
+
+}