Skip to content

Commit

Permalink
Merge branch 'eriknovak:main' into main
Browse files Browse the repository at this point in the history
  • Loading branch information
eriknovak authored Dec 18, 2024
2 parents c652040 + 7599267 commit c97aa85
Showing 1 changed file with 2 additions and 16 deletions.
18 changes: 2 additions & 16 deletions anonipy/anonymize/generators/llm_label_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ def _generate_response(

# tokenize the message
input_ids = self.tokenizer.apply_chat_template(
message, tokenize=True, return_tensors="pt"
message, tokenize=True, return_tensors="pt", add_generation_prompt=True
).to(self.model.device)

# generate the response
Expand All @@ -166,18 +166,4 @@ def _generate_response(
response = self.tokenizer.decode(
output_ids[0][len(input_ids[0]) :], skip_special_tokens=True
)
return self._parse_response(response)

def _parse_response(self, response: str) -> str:
"""Parse the response from the LLM.
Args:
response: The response to parse.
Returns:
The parsed response.
"""

match = re.search(r"assistant\s*(.*)", response, re.IGNORECASE | re.DOTALL)
return match.group(1).strip() if match else response
return response

0 comments on commit c97aa85

Please sign in to comment.