Add conditions and restraints to LLM prompt

eriknovak · Nov 9, 2024 · 25a6e44 · 25a6e44
1 parent 6a3ffb2
commit 25a6e44
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/anonipy/anonymize/generators/llm_label_generator.py b/anonipy/anonymize/generators/llm_label_generator.py
@@ -98,7 +98,7 @@ def generate(
             },
             {
                 "role": "user",
-                "content": f"What is a random {add_entity_attrs} {entity.label} replacement for {entity.text}? Respond only with the replacement.",
+                "content": f"What is a random {add_entity_attrs} {entity.label} replacement for {entity.text}? Respond only with the replacement, on additional information.",
             },
         ]
         return self._generate_response(message, temperature, top_p)
@@ -179,5 +179,9 @@ def _parse_response(self, response: str) -> str:
 
         """
 
-        match = re.search(r"assistant\s*(.*)", response, re.IGNORECASE | re.DOTALL)
-        return match.group(1).strip() if match else response
+        res_match = re.search(r"assistant\s*(.*)", response, re.IGNORECASE | re.DOTALL)
+        tmp_response = res_match.group(1).strip() if res_match else response
+
+        res_match = re.search(r"\:\s*(.*)", tmp_response, re.IGNORECASE | re.DOTALL)
+
+        return