From 073e991ad42157d78a089b23ebe3da0bf1da7907 Mon Sep 17 00:00:00 2001 From: Zdenek Kasner Date: Mon, 9 Dec 2024 17:05:27 +0100 Subject: [PATCH] Add stopping sequence parameter for ollama --- factgenie/models.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/factgenie/models.py b/factgenie/models.py index f0d026d..e4dd73b 100644 --- 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