refactor: improve auto-model resolution for task

Author: CodeWithKyrian

src/Pipelines/Pipeline.php

@@ -73,7 +73,25 @@ function pipeline(
 
     $modelName ??= $task->defaultModelName();
 
-    $model = $task->autoModel($modelName, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress);
+    $modelClass = $task->autoModelClass();
+
+    $modelClass = is_array($modelClass) ? $modelClass : [$modelClass];
+
+    $model = null;
+    $lastException = null;
+
+    foreach ($modelClass as $modelClass) {
+        try {
+            $model = $modelClass::fromPretrained($modelName, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress);
+            break;
+        } catch (\Throwable $e) {
+            $lastException = $e;
+        }
+    }
+
+    if ($model === null) {
+        throw new \RuntimeException('Could not instantiate model for task: ' . $task->value, 0, $lastException);
+    }
 
     $tokenizer = $task->autoTokenizer($modelName, $cacheDir, $revision, $onProgress);
 

src/Pipelines/Task.php

@@ -151,61 +151,50 @@ public function defaultModelName(): string
         };
     }
 
-    public function autoModel(
-        string    $modelNameOrPath,
-        bool      $quantized = true,
-        ?array    $config = null,
-        ?string   $cacheDir = null,
-        string    $revision = 'main',
-        ?string   $modelFilename = null,
-        ?callable $onProgress = null
-    ): PretrainedModel
+    /**
+     * @return class-string<PretrainedModel>|array<class-string<PretrainedModel>>
+     */
+    public function autoModelClass(): string|array
     {
         return match ($this) {
             self::SentimentAnalysis,
             self::TextClassification,
-            self::ZeroShotClassification => AutoModelForSequenceClassification::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ZeroShotClassification => AutoModelForSequenceClassification::class,
 
-            self::FillMask => AutoModelForMaskedLM::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::FillMask => AutoModelForMaskedLM::class,
 
-            self::QuestionAnswering => AutoModelForQuestionAnswering::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::QuestionAnswering => AutoModelForQuestionAnswering::class,
 
             self::FeatureExtraction,
-            self::Embeddings => AutoModel::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::Embeddings => AutoModel::class,
 
             self::Text2TextGeneration,
             self::Translation,
-            self::Summarization => AutoModelForSeq2SeqLM::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::Summarization => AutoModelForSeq2SeqLM::class,
 
-            self::TextGeneration => AutoModelForCausalLM::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::TextGeneration => AutoModelForCausalLM::class,
 
             self::TokenClassification,
-            self::Ner => AutoModelForTokenClassification::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::Ner => AutoModelForTokenClassification::class,
 
-            self::ImageToText => AutoModelForVision2Seq::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ImageToText => AutoModelForVision2Seq::class,
 
-            self::ImageClassification => AutoModelForImageClassification::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ImageClassification => AutoModelForImageClassification::class,
 
-            self::ImageFeatureExtraction => AutoModelForImageFeatureExtraction::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ImageFeatureExtraction => [AutoModelForImageFeatureExtraction::class, AutoModel::class],
 
-            self::ZeroShotImageClassification => AutoModel::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ZeroShotImageClassification => AutoModel::class,
 
-            self::ImageToImage => AutoModelForImageToImage::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ImageToImage => AutoModelForImageToImage::class,
 
-            self::ObjectDetection => AutoModelForObjectDetection::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ObjectDetection => AutoModelForObjectDetection::class,
 
-            self::ZeroShotObjectDetection => AutoModelForZeroShotObjectDetection::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::ZeroShotObjectDetection => AutoModelForZeroShotObjectDetection::class,
 
-            self::AudioClassification => AutoModelForAudioClassification::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress),
+            self::AudioClassification => AutoModelForAudioClassification::class,
 
             self::ASR,
-            self::AutomaticSpeechRecognition => (function () use ($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress) {
-                try {
-                    return AutoModelForSpeechSeq2Seq::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress);
-                } catch (UnsupportedModelTypeException) {
-                    return AutoModelForCTC::fromPretrained($modelNameOrPath, $quantized, $config, $cacheDir, $revision, $modelFilename, $onProgress);
-                }
-            })(),
+            self::AutomaticSpeechRecognition => [AutoModelForSpeechSeq2Seq::class, AutoModelForCTC::class],
         };
     }
 
@@ -214,8 +203,7 @@ public function autoTokenizer(
         ?string   $cacheDir = null,
         string    $revision = 'main',
         ?callable $onProgress = null
-    ): ?PreTrainedTokenizer
-    {
+    ): ?PreTrainedTokenizer {
         return match ($this) {
 
             self::ImageClassification,
@@ -252,8 +240,7 @@ public function autoProcessor(
         ?string   $cacheDir = null,
         string    $revision = 'main',
         ?callable $onProgress = null
-    ): ?Processor
-    {
+    ): ?Processor {
         return match ($this) {
 
             self::ImageToText,