Merge pull request #9 from MinuraPunchihewa/development

[PR] dev to main: updated the unit tests for the text generation tasks
MinuraPunchihewa · Apr 29, 2023 · 0d819cb · 0d819cb
2 parents 0e8be89 + b7ff61b
commit 0d819cb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 19 deletions.
diff --git a/tests/nlp/test_nlp.py b/tests/nlp/test_nlp.py
@@ -144,16 +144,8 @@ def test_text_generation(self):
         text = "The answer to the universe is"
 
         try:
-            self.assertEqual(
-                self.nlp.text_generation(text),
-                [
-                    {'generated_text': 'The answer to the universe is simple: The answer can be '
-                                    'proved. The first thing you should do is make an '
-                                    'appointment with a psychic called S.K. Gupta, a.k.a. '
-                                    'Shiva (his real name) who has the most'
-                     }
-                ]
-            )
+            prediction = self.nlp.text_generation(text)
+            self.assertTrue(prediction[0]['generated_text'].startswith(text))
         except HTTPServiceUnavailableException:
             self.assertRaises(HTTPServiceUnavailableException, lambda: self.nlp.text_generation(text))
 

diff --git a/tests/nlp/test_nlp_in_df.py b/tests/nlp/test_nlp_in_df.py
@@ -116,15 +116,9 @@ def test_text_generation_in_df(self):
         df = pd.DataFrame(texts, columns=['texts'])
 
         try:
-            assert_frame_equal(
-                self.nlp.text_generation_in_df(df, 'texts'),
-                pd.DataFrame(
-                    {
-                        "texts": texts,
-                        "predictions": ["The answer to the universe is not a random, \"good and bad.\" At least not yet.\n\nIf you're the kind of guy who enjoys reading at least one novel an hour, you would better be getting into it. But you're"],
-                    }
-                ),
-            )
+            predictions_df = self.nlp.text_generation_in_df(df, 'texts')
+            for index, row in predictions_df.iterrows():
+                self.assertTrue(row['predictions'].startswith(row['texts']))
         except HTTPServiceUnavailableException:
             self.assertRaises(HTTPServiceUnavailableException, lambda: self.nlp.text_generation_in_df(df, 'texts'))