fix grok-1 tests

Signed-off-by: Mamta Singh <[email protected]>

Author: mamtsing

QEfficient/utils/test_utils.py

@@ -165,7 +165,34 @@ class ModelConfig:
     }
 
     EXTERNAL_MODELS = {
-        "hpcai-tech/grok-1",
+        "hpcai-tech/grok-1": {
+            "pytorch_hf_tokens": [
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+                391,
+            ]
+        }
     }
 
     SWIFTKV_MODELS = {

pyproject.toml

@@ -39,11 +39,11 @@ dependencies = [
     "fire",
     "py7zr",
     "torchmetrics==1.7.0",
-    "torch==2.7.1; platform_machine=='aarch64'",
+    "torch==2.7.0; platform_machine=='aarch64'",
     # Specifying torch cpu package URL per python version, update the list once pytorch releases whl for python>3.11
     "torch@https://download.pytorch.org/whl/cpu/torch-2.4.1%2Bcpu-cp38-cp38-linux_x86_64.whl ; python_version=='3.8' and platform_machine=='x86_64'",
-    "torch@https://download.pytorch.org/whl/cpu/torch-2.7.1%2Bcpu-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_machine=='x86_64'",
-    "torch@https://download.pytorch.org/whl/cpu/torch-2.7.1%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_machine=='x86_64'",
+    "torch@https://download.pytorch.org/whl/cpu/torch-2.7.0%2Bcpu-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_machine=='x86_64'",
+    "torch@https://download.pytorch.org/whl/cpu/torch-2.7.0%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_machine=='x86_64'",
 ]
 
 [project.optional-dependencies]

tests/transformers/models/test_causal_lm_models.py

@@ -170,8 +170,10 @@ def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(
         Constants.CTX_LEN,
     )
 
-    if model_name not in ModelConfig.SWIFTKV_MODELS:
+    if model_name not in (ModelConfig.SWIFTKV_MODELS and ModelConfig.EXTERNAL_MODELS):
         pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch(model_hf)
+    if model_name in ModelConfig.EXTERNAL_MODELS:
+        pytorch_hf_tokens = ModelConfig.EXTERNAL_MODELS[model_name]["pytorch_hf_tokens"]
 
     is_tlm = False if num_speculative_tokens is None else True
     qeff_model = QEFFAutoModelForCausalLM(
@@ -232,10 +234,13 @@ def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(
         full_batch_size,
     )
 
-    if model_name not in ModelConfig.SWIFTKV_MODELS:
+    if model_name not in (ModelConfig.SWIFTKV_MODELS and ModelConfig.EXTERNAL_MODELS):
         pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch_CB(model_hf)
         pytorch_hf_tokens = np.vstack(pytorch_hf_tokens)
 
+    if model_name in ModelConfig.EXTERNAL_MODELS:
+        pytorch_hf_tokens = [pytorch_hf_tokens for _ in range(full_batch_size)]
+
     qeff_model = QEFFAutoModelForCausalLM(
         model_hf, continuous_batching=True, is_tlm=is_tlm, pretrained_model_name_or_path=model_name
     )