diff --git a/openlrc/translate.py b/openlrc/translate.py index 431ce34..8a72cd9 100644 --- a/openlrc/translate.py +++ b/openlrc/translate.py @@ -21,8 +21,8 @@ class Translator(ABC): @abstractmethod - def translate(self, texts: Union[str, List[str]], src_lang: str, target_lang: str, - info: TranslateInfo) -> List[str]: + def translate(self, texts: Union[str, List[str]], src_lang: str, target_lang: str, info: TranslateInfo) \ + -> List[str]: pass @@ -44,6 +44,7 @@ def __init__(self, chatbot_model: str = 'gpt-3.5-turbo', fee_limit: float = 0.3, self.api_fee = 0 self.intercept_line = intercept_line self.retry_model = retry_model + self.use_retry_cnt = 0 @staticmethod def list_chatbots() -> List[str]: @@ -82,26 +83,30 @@ def make_chunks(texts: List[str], chunk_size: int = 30) -> List[List[Tuple[int, def _translate_chunk(self, translator_agent: ChunkedTranslatorAgent, chunk: List[Tuple[int, str]], context: TranslationContext, chunk_id: int, - retry_agent: Optional[ChunkedTranslatorAgent] = None) -> Tuple[ - List[str], TranslationContext]: + retry_agent: Optional[ChunkedTranslatorAgent] = None) -> Tuple[List[str], TranslationContext]: """ Translate a single chunk of text. """ - translated, context = translator_agent.translate_chunk(chunk_id, chunk, context) - - if len(translated) != len(chunk) and translator_agent.info.glossary: - logger.warning(f'Cannot translate chunk {chunk_id} with glossary, trying to remove glossary.') - translated, context = translator_agent.translate_chunk(chunk_id, chunk, context, use_glossary=False) - - if retry_agent and len(translated) != len(chunk): - logger.warning( - f'Trying to change chatbot to keep performing chunked translation. Retry chatbot: {retry_agent}') - # TODO: use the retry_agent for all the next chunks - translated, context = retry_agent.translate_chunk(chunk_id, chunk, context) - if len(translated) != len(chunk) and retry_agent.info.glossary: - logger.warning(f'New bot: Trying to remove glossary to keep performing chunked translation.') - translated, context = retry_agent.translate_chunk(chunk_id, chunk, context, use_glossary=False) + def handle_translation(agent: ChunkedTranslatorAgent) -> Tuple[List[str], TranslationContext]: + trans, updated_context = agent.translate_chunk(chunk_id, chunk, context) + if len(trans) != len(chunk) and agent.info.glossary: + logger.warning(f'Agent {agent}: Removing glossary for chunk {chunk_id}.') + trans, updated_context = agent.translate_chunk(chunk_id, chunk, context, use_glossary=False) + return trans, updated_context + + if self.use_retry_cnt == 0 or not retry_agent: + translated, context = handle_translation(translator_agent) + + if retry_agent and len(translated) != len(chunk): + self.use_retry_cnt = 3 # Use retry_agent for the next 3 chunks + logger.warning( + f'Using retry agent {retry_agent} for chunk {chunk_id}, and next {self.use_retry_cnt} chunks.') + translated, context = handle_translation(retry_agent) + else: + logger.info(f'Using retry agent for chunk {chunk_id}, remaining retries: {self.use_retry_cnt}') + translated, context = handle_translation(retry_agent) + self.use_retry_cnt -= 1 return translated, context