Workaround for minilmv2 pytorch output issue

groq · Feb 1, 2024 · a4c8e47 · a4c8e47
1 parent 2582343
commit a4c8e47
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/demo_helpers/demo_helpers/compute_performance.py b/demo_helpers/demo_helpers/compute_performance.py
@@ -182,8 +182,11 @@ def pytorch_model_inference(dataset, model):
             out = model(**inputs)
 
             if not isinstance(out, torch.Tensor):
-                if isinstance(out, tuple) and len(out) == 1:
-                    out = out[0]
+                if isinstance(out, tuple):
+                    if len(out) == 1:
+                        out = out[0]
+                    else:
+                        raise ValueError("Cannot handle tuple with len", len(out))
                 elif isinstance(out, dict):
                     if "logits" in out:
                         out = out.logits

diff --git a/proof_points/natural_language_processing/minilm/minilmv2.py b/proof_points/natural_language_processing/minilm/minilmv2.py
@@ -19,9 +19,7 @@ def evaluate_minilm(rebuild_policy=None, should_execute=True):
 
     # load pre-trained torch model
     tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-    model = AutoModel.from_pretrained(
-        "sentence-transformers/all-MiniLM-L6-v2", torchscript=True
-    )
+    model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
 
     # dummy inputs to generate the groq model
     max_seq_length = 128