fixed tests

amaiya · Jun 14, 2024 · 0d39e51 · 0d39e51
1 parent de5f383
commit 0d39e51
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tests/test_ner.py b/tests/test_ner.py
@@ -25,7 +25,7 @@ def setUp(self):
 
     def test_ner(self):
         model = txt.sequence_tagger(
-            "bilstm-transformer", self.preproc, transformer_model="bert-base-cased"
+            "bilstm-transformer", self.preproc, transformer_model="roberta-base"
         )
         learner = ktrain.get_learner(
             model, train_data=self.trn, val_data=self.val, batch_size=128

diff --git a/tests/test_qa.py b/tests/test_qa.py
@@ -41,7 +41,7 @@ def test_qa(self):
         top_answer = answers[0]["answer"]
         self.assertEqual(top_answer, "in october of 1997")
 
-    @skip("temporarily disabled")
+    # @skip("temporarily disabled")
     def test_extractor(self):
 
         # data = ['Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .',
@@ -76,7 +76,8 @@ def test_extractor(self):
         ]
         from ktrain.text import AnswerExtractor
 
-        ae = AnswerExtractor(framework="pt", device="cpu", quantize=True)
+        # ae = AnswerExtractor(framework="pt", device="cpu", quantize=True)
+        ae = AnswerExtractor()
         import pandas as pd
 
         pd.set_option("display.max_colwidth", None)