fixed: Unbound Error and Key Error.

JohnSnowLabs · Sep 3, 2024 · 258a0f7 · 258a0f7
1 parent 16fee46
commit 258a0f7
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 23 deletions.
diff --git a/langtest/datahandler/datasource.py b/langtest/datahandler/datasource.py
@@ -957,7 +957,7 @@ def _import_data(self, file_name, **kwargs) -> List[Sample]:
                 import ast
 
                 i["transformations"] = ast.literal_eval(temp)
-            else:
+            elif "transformations" in i:
                 i.pop("transformations")
             sample = self.task.get_sample_class(**i)
             samples.append(sample)

diff --git a/langtest/modelhandler/jsl_modelhandler.py b/langtest/modelhandler/jsl_modelhandler.py
@@ -450,31 +450,25 @@ def predict(
             SequenceClassificationOutput: Classification output from SparkNLP LightPipeline.
         """
         prediction_metadata = self.model.fullAnnotate(text)[0][self.output_col]
+        prediction = []
+
+        if len(prediction_metadata) > 0:
+            prediction_metadata = prediction_metadata[0].metadata
+            prediction = [
+                {"label": x, "score": y} for x, y in prediction_metadata.items()
+            ]
 
         if self.multi_label_classifier:
-            multi_label = True
-            if len(prediction_metadata) > 0:
-                prediction_metadata = prediction_metadata[0].metadata
-
-                prediction = [
-                    {"label": x, "score": y} for x, y in prediction_metadata.items()
-                ]
-                # filter based on the threshold value with score greater than threshold
-                prediction = [x for x in prediction if float(x["score"]) > self.threshold]
-
-                return SequenceClassificationOutput(
-                    text=text,
-                    predictions=prediction,
-                    multi_label=multi_label,
-                )
-            else:
-                return SequenceClassificationOutput(
-                    text=text, predictions=[], multi_label=multi_label
-                )
+            prediction = [x for x in prediction if float(x["score"]) > self.threshold]
 
-        else:
-            if not return_all_scores:
-                prediction = [max(prediction, key=lambda x: x["score"])]
+            return SequenceClassificationOutput(
+                text=text,
+                predictions=prediction,
+                multi_label=self.multi_label_classifier,
+            )
+
+        if not return_all_scores:
+            prediction = [max(prediction, key=lambda x: x["score"])]
 
         return SequenceClassificationOutput(text=text, predictions=prediction)