Revert "Add conditions and restraints to LLM prompt"

This reverts commit 25a6e44.
eriknovak · Nov 9, 2024 · 63a5589 · 63a5589
1 parent 77307a3
commit 63a5589
Showing 1 changed file with 3 additions and 7 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, on additional information.",
+                "content": f"What is a random {add_entity_attrs} {entity.label} replacement for {entity.text}? Respond only with the replacement.",
             },
         ]
         return self._generate_response(message, temperature, top_p)
@@ -179,9 +179,5 @@ def _parse_response(self, response: str) -> str:
 
         """
 
-        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
+        match = re.search(r"assistant\s*(.*)", response, re.IGNORECASE | re.DOTALL)
+        return match.group(1).strip() if match else response