option to select claude model

This commit is contained in:
cce 2024-11-27 21:35:24 -05:00
parent daea974d68
commit efaa673226
No known key found for this signature in database
3 changed files with 11 additions and 1 deletions

View File

@ -375,7 +375,7 @@ So you are close to reaching the limit. You have to choose your own value, there
API_KEY = options.deepl_key or env.get("BBM_DEEPL_API_KEY")
if not API_KEY:
raise Exception("Please provide deepl key")
elif options.model == "claude":
elif options.model.startswith("claude"):
API_KEY = options.claude_key or env.get("BBM_CLAUDE_API_KEY")
if not API_KEY:
raise Exception("Please provide claude key")
@ -494,6 +494,8 @@ So you are close to reaching the limit. You have to choose your own value, there
e.translate_model.set_gpt4omini_models()
if options.model == "gpt4o":
e.translate_model.set_gpt4o_models()
if options.model.startswith("claude-"):
e.translate_model.set_claude_model(options.model)
if options.block_size > 0:
e.block_size = options.block_size
if options.batch_flag:

View File

@ -21,6 +21,11 @@ MODEL_DICT = {
"deepl": DeepL,
"deeplfree": DeepLFree,
"claude": Claude,
"claude-3-5-sonnet-latest": Claude,
"claude-3-5-sonnet-20241022": Claude,
"claude-3-5-sonnet-20240620": Claude,
"claude-3-5-haiku-latest": Claude,
"claude-3-5-haiku-20241022": Claude,
"gemini": Gemini,
"geminipro": Gemini,
"groq": GroqClient,

View File

@ -30,6 +30,9 @@ class Claude(Base):
def rotate_key(self):
pass
def set_claude_model(self, model_name):
self.model = model_name
def translate(self, text):
print(text)
self.rotate_key()