Add external data model architecture tests

Author: xenova

tests/models.test.js

@@ -2,7 +2,24 @@
  * Test that models loaded outside of the `pipeline` function work correctly (e.g., `AutoModel.from_pretrained(...)`);
  */
 
-import { AutoTokenizer, AutoModel, BertModel, GPT2Model, T5ForConditionalGeneration, BertTokenizer, GPT2Tokenizer, T5Tokenizer } from "../src/transformers.js";
+import {
+  AutoTokenizer,
+  AutoProcessor,
+  BertForMaskedLM,
+  GPT2LMHeadModel,
+  T5ForConditionalGeneration,
+  BertTokenizer,
+  GPT2Tokenizer,
+  T5Tokenizer,
+  LlamaTokenizer,
+  LlamaForCausalLM,
+  WhisperForConditionalGeneration,
+  WhisperProcessor,
+  AutoModelForMaskedLM,
+  AutoModelForCausalLM,
+  AutoModelForSeq2SeqLM,
+  AutoModelForSpeechSeq2Seq,
+} from "../src/transformers.js";
 import { init, MAX_TEST_EXECUTION_TIME, DEFAULT_MODEL_OPTIONS } from "./init.js";
 import { compare, collect_and_execute_tests } from "./test_utils.js";
 
@@ -12,44 +29,51 @@ init();
 describe("Loading different architecture types", () => {
   // List all models which will be tested
   const models_to_test = [
-    // [name, modelClass, tokenizerClass]
-    ["hf-internal-testing/tiny-random-BertForMaskedLM", BertModel, BertTokenizer], // Encoder-only
-    ["hf-internal-testing/tiny-random-GPT2LMHeadModel", GPT2Model, GPT2Tokenizer], // Decoder-only
-    ["hf-internal-testing/tiny-random-T5ForConditionalGeneration", T5ForConditionalGeneration, T5Tokenizer], // Encoder-decoder
+    // [name, [AutoModelClass, ModelClass], [AutoProcessorClass, ProcessorClass], [modelOptions?], [modality?]]
+    ["hf-internal-testing/tiny-random-BertForMaskedLM", [AutoModelForMaskedLM, BertForMaskedLM], [AutoTokenizer, BertTokenizer]], // Encoder-only
+    ["hf-internal-testing/tiny-random-GPT2LMHeadModel", [AutoModelForCausalLM, GPT2LMHeadModel], [AutoTokenizer, GPT2Tokenizer]], // Decoder-only
+    ["hf-internal-testing/tiny-random-T5ForConditionalGeneration", [AutoModelForSeq2SeqLM, T5ForConditionalGeneration], [AutoTokenizer, T5Tokenizer]], // Encoder-decoder
+    ["onnx-internal-testing/tiny-random-LlamaForCausalLM-ONNX_external", [AutoModelForCausalLM, LlamaForCausalLM], [AutoTokenizer, LlamaTokenizer]], // Decoder-only w/ external data
+    ["onnx-internal-testing/tiny-random-WhisperForConditionalGeneration-ONNX_external", [AutoModelForSpeechSeq2Seq, WhisperForConditionalGeneration], [AutoProcessor, WhisperProcessor], {}, "audio"], // Encoder-decoder-only w/ external data
   ];
 
   const texts = ["Once upon a time", "I like to eat apples"];
 
-  for (const [model_id, modelClass, tokenizerClass] of models_to_test) {
+  for (const [model_id, models, processors, modelOptions, modality] of models_to_test) {
     // Test that both the auto model and the specific model work
-    const tokenizers = [AutoTokenizer, tokenizerClass];
-    const models = [AutoModel, modelClass];
-
-    for (let i = 0; i < tokenizers.length; ++i) {
-      const tokenizerClassToTest = tokenizers[i];
+    for (let i = 0; i < processors.length; ++i) {
+      const processorClassToTest = processors[i];
       const modelClassToTest = models[i];
 
       it(
         `${model_id} (${modelClassToTest.name})`,
         async () => {
-          // Load model and tokenizer
-          const tokenizer = await tokenizerClassToTest.from_pretrained(model_id);
-          const model = await modelClassToTest.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
+          // Load model and processor
+          const processor = await processorClassToTest.from_pretrained(model_id);
+          const model = await modelClassToTest.from_pretrained(model_id, modelOptions ?? DEFAULT_MODEL_OPTIONS);
+
+          const tests = modality === "audio"
+            ? [new Float32Array(16000)]
+            : [
+              texts[0], // single
+              texts, // batched
+            ]
 
-          const tests = [
-            texts[0], // single
-            texts, // batched
-          ];
           for (const test of tests) {
-            const inputs = await tokenizer(test, { truncation: true, padding: true });
+            const inputs = await processor(test, { truncation: true, padding: true });
             if (model.config.is_encoder_decoder) {
-              inputs.decoder_input_ids = inputs.input_ids;
+              if (model.config.model_type === "whisper") {
+                inputs.decoder_input_ids = processor.tokenizer(texts[0]).input_ids;
+              } else {
+                inputs.decoder_input_ids = inputs.input_ids;
+              }
             }
             const output = await model(inputs);
 
             if (output.logits) {
               // Ensure correct shapes
-              const expected_shape = [...inputs.input_ids.dims, model.config.vocab_size];
+              const input_ids = inputs.input_ids ?? inputs.decoder_input_ids;
+              const expected_shape = [...input_ids.dims, model.config.vocab_size];
               const actual_shape = output.logits.dims;
               compare(expected_shape, actual_shape);
             } else if (output.last_hidden_state) {