Update HallucinationScorer to accept list of contexts

wandb · Feb 5, 2025 · c97963d · c97963d
1 parent dd13e86
commit c97963d
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/weave/scorers/hallucination_scorer.py b/weave/scorers/hallucination_scorer.py
@@ -306,15 +306,17 @@ def load_model(self) -> None:
         self.model.eval()
 
     @weave.op
-    def score(self, query: str, context: str, output: str) -> dict[str, Any]:
+    def score(self, query: str, context: str | list[str], output: str) -> dict[str, Any]:
         if self.import_failed:
             return {
                 "flagged": False,
                 "extras": {
                     "error": "Unable to import required libraries. Please install transformers."
                 },
             }
-        inps = query + "\n\n" + context
+        # Handle context being either string or list of strings
+        context_str = "\n\n".join(context) if isinstance(context, list) else context
+        inps = query + "\n\n" + context_str
         outs = output
 
         inps_toks = self._tokenizer(inps, truncation=False)