mirror of
https://github.com/lightbroker/llmsecops-research.git
synced 2026-05-13 04:44:45 +02:00
Adam Wilson
22 lines
969 B
Python
22 lines
969 B
Python
from src.text_generation.common.constants import Constants
|
|
from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService
|
|
|
|
|
|
class ResponseProcessingService(
|
|
AbstractResponseProcessingService):
|
|
|
|
def __init__(self):
|
|
self.constants = Constants()
|
|
|
|
def process_text_generation_output(self, raw_output: str) -> str:
|
|
if self.constants.PHI_3_ASSISTANT_START_TOKEN in raw_output:
|
|
# split at assistant token and take everything after it
|
|
parts = raw_output.split(self.constants.PHI_3_ASSISTANT_START_TOKEN)
|
|
answer = parts[-1].strip()
|
|
# remove trailing <|end|> tokens if present
|
|
if answer.endswith(self.constants.PHI_3_END_TOKEN):
|
|
answer = answer[:-(len(self.constants.PHI_3_END_TOKEN))].strip()
|
|
return answer
|
|
else:
|
|
# return raw original (fallback)
|
|
return raw_output.strip() |