Add stopping sequence parameter for ollama

ufal · Dec 9, 2024 · 073e991 · 073e991
1 parent 9ca7b44
commit 073e991
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/factgenie/models.py b/factgenie/models.py
@@ -413,12 +413,20 @@ def get_optional_fields(self):
         }
 
     def postprocess_output(self, output):
-        if self.config.get("extra_args", {}).get("remove_suffix", ""):
+        extra_args = self.config.get("extra_args", {})
+
+        if extra_args.get("remove_suffix", ""):
             suffix = self.config["extra_args"]["remove_suffix"]
 
             if output.endswith(suffix):
                 output = output[: -len(suffix)]
 
+        if extra_args.get("stopping_sequence", False):
+            stopping_sequence = self.config["extra_args"]["stopping_sequence"]
+
+            if stopping_sequence in output:
+                output = output[: output.index(stopping_sequence)]
+
         output = output.strip()
         return output