Update binoculars/detector.py

binoculars/detector.py

@@ -99,10 +99,10 @@ class Binoculars(object):
         binoculars_scores = binoculars_scores.tolist()
         return binoculars_scores[0] if isinstance(input_text, str) else binoculars_scores
 
-    def predict(self, input_text: Union[list[str], str]) -> Union[list[str], str]:
+    def predict(self, input_text: Union[list[str], str]) -> Union[str, list[str]]:
         binoculars_scores = np.array(self.compute_score(input_text))
-        pred = np.where(binoculars_scores < self.threshold,
-                        "Most likely AI-generated",
-                        "Most likely human-generated"
-                        ).tolist()
-        return pred[0] if isinstance(input_text, str) else pred
+        
+        if isinstance(input_text, str):
+            return "Most likely AI-generated" if binoculars_scores < self.threshold else "Most likely human-generated"
+        else:
+            return ["Most likely AI-generated" if score < self.threshold else "Most likely human-generated" for score in binoculars_scores]