mirror of
https://github.com/yihong0618/bilingual_book_maker.git
synced 2025-06-06 11:35:49 +00:00
b80f1ba785
* option to select claude model * add context_flag and context_paragraph_limit option for claude * reformat with black * remove nonexistent model
112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
import re
|
|
import time
|
|
from rich import print
|
|
from anthropic import Anthropic
|
|
|
|
from .base_translator import Base
|
|
|
|
|
|
class Claude(Base):
|
|
def __init__(
|
|
self,
|
|
key,
|
|
language,
|
|
api_base=None,
|
|
prompt_template=None,
|
|
prompt_sys_msg=None,
|
|
temperature=1.0,
|
|
context_flag=False,
|
|
context_paragraph_limit=5,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__(key, language)
|
|
self.api_url = api_base or "https://api.anthropic.com"
|
|
self.client = Anthropic(base_url=api_base, api_key=key, timeout=20)
|
|
self.model = "claude-3-5-sonnet-20241022" # default it for now
|
|
self.language = language
|
|
self.prompt_template = (
|
|
prompt_template
|
|
or "Help me translate the text within triple backticks into {language} and provide only the translated result.\n```{text}```"
|
|
)
|
|
self.prompt_sys_msg = prompt_sys_msg or ""
|
|
self.temperature = temperature
|
|
self.context_flag = context_flag
|
|
self.context_list = []
|
|
self.context_translated_list = []
|
|
self.context_paragraph_limit = context_paragraph_limit
|
|
|
|
def rotate_key(self):
|
|
pass
|
|
|
|
def set_claude_model(self, model_name):
|
|
self.model = model_name
|
|
|
|
def create_messages(self, text, intermediate_messages=None):
|
|
"""Create messages for the current translation request"""
|
|
current_msg = {
|
|
"role": "user",
|
|
"content": self.prompt_template.format(
|
|
text=text,
|
|
language=self.language,
|
|
),
|
|
}
|
|
|
|
messages = []
|
|
if intermediate_messages:
|
|
messages.extend(intermediate_messages)
|
|
messages.append(current_msg)
|
|
|
|
return messages
|
|
|
|
def create_context_messages(self):
|
|
"""Create a message pair containing all context paragraphs"""
|
|
if not self.context_flag or not self.context_list:
|
|
return []
|
|
|
|
# Create a single message pair for all previous context
|
|
return [
|
|
{
|
|
"role": "user",
|
|
"content": self.prompt_template.format(
|
|
text="\n\n".join(self.context_list),
|
|
language=self.language,
|
|
),
|
|
},
|
|
{"role": "assistant", "content": "\n\n".join(self.context_translated_list)},
|
|
]
|
|
|
|
def save_context(self, text, t_text):
|
|
"""Save the current translation pair to context"""
|
|
if not self.context_flag:
|
|
return
|
|
|
|
self.context_list.append(text)
|
|
self.context_translated_list.append(t_text)
|
|
|
|
# Keep only the most recent paragraphs within the limit
|
|
if len(self.context_list) > self.context_paragraph_limit:
|
|
self.context_list.pop(0)
|
|
self.context_translated_list.pop(0)
|
|
|
|
def translate(self, text):
|
|
print(text)
|
|
self.rotate_key()
|
|
|
|
# Create messages with context
|
|
messages = self.create_messages(text, self.create_context_messages())
|
|
|
|
r = self.client.messages.create(
|
|
max_tokens=4096,
|
|
messages=messages,
|
|
system=self.prompt_sys_msg,
|
|
temperature=self.temperature,
|
|
model=self.model,
|
|
)
|
|
t_text = r.content[0].text
|
|
|
|
if self.context_flag:
|
|
self.save_context(text, t_text)
|
|
|
|
print("[bold green]" + re.sub("\n{3,}", "\n\n", t_text) + "[/bold green]")
|
|
return t_text
|