diff --git a/src-python/config.py b/src-python/config.py index aca3203c..7fec780f 100644 --- a/src-python/config.py +++ b/src-python/config.py @@ -9,6 +9,7 @@ import torch from device_manager import device_manager from models.translation.translation_languages import translation_lang from models.translation.translation_utils import ctranslate2_weights +from models.translation.translation_plamo import _MODELS as plamo_models from models.transcription.transcription_languages import transcription_lang from models.transcription.transcription_whisper import _MODELS as whisper_models from utils import errorLogging, validateDictStructure @@ -119,6 +120,10 @@ class Config: def SELECTABLE_TRANSLATION_ENGINE_LIST(self): return self._SELECTABLE_TRANSLATION_ENGINE_LIST + @property + def SELECTABLE_PLAMO_MODEL_LIST(self): + return self._SELECTABLE_PLAMO_MODEL_LIST + @property def SELECTABLE_TRANSCRIPTION_ENGINE_LIST(self): return self._SELECTABLE_TRANSCRIPTION_ENGINE_LIST @@ -826,6 +831,18 @@ class Config: self._WHISPER_WEIGHT_TYPE = value self.saveConfig(inspect.currentframe().f_code.co_name, value) + @property + @json_serializable('PLAMO_MODEL') + def PLAMO_MODEL(self): + return self._PLAMO_MODEL + + @PLAMO_MODEL.setter + def PLAMO_MODEL(self, value): + if isinstance(value, str): + if value in self.SELECTABLE_PLAMO_MODEL_LIST: + self._PLAMO_MODEL = value + self.saveConfig(inspect.currentframe().f_code.co_name, value) + @property @json_serializable('AUTO_CLEAR_MESSAGE_BOX') def AUTO_CLEAR_MESSAGE_BOX(self): @@ -1043,6 +1060,7 @@ class Config: self._SELECTABLE_CTRANSLATE2_WEIGHT_TYPE_LIST = ctranslate2_weights.keys() self._SELECTABLE_WHISPER_WEIGHT_TYPE_LIST = whisper_models.keys() self._SELECTABLE_TRANSLATION_ENGINE_LIST = translation_lang.keys() + self._SELECTABLE_PLAMO_MODEL_LIST = plamo_models self._SELECTABLE_TRANSCRIPTION_ENGINE_LIST = list(transcription_lang[list(transcription_lang.keys())[0]].values())[0].keys() self._SELECTABLE_UI_LANGUAGE_LIST = ["en", "ja", "ko", "zh-Hant", "zh-Hans"] self._COMPUTE_MODE = "cuda" if torch.cuda.is_available() else "cpu" @@ -1191,6 +1209,7 @@ class Config: self._SELECTED_TRANSCRIPTION_COMPUTE_DEVICE = copy.deepcopy(self.SELECTABLE_COMPUTE_DEVICE_LIST[0]) self._CTRANSLATE2_WEIGHT_TYPE = "m2m100_418M-ct2-int8" self._WHISPER_WEIGHT_TYPE = "base" + self._PLAMO_MODEL = "plamo-2.0-prime" self._AUTO_CLEAR_MESSAGE_BOX = True self._SEND_ONLY_TRANSLATED_MESSAGES = False self._OVERLAY_SMALL_LOG = False diff --git a/src-python/controller.py b/src-python/controller.py index 514a105e..dc3418d5 100644 --- a/src-python/controller.py +++ b/src-python/controller.py @@ -1431,6 +1431,36 @@ class Controller: self.updateTranslationEngineAndEngineList() return {"status":200, "result":config.AUTH_KEYS[translator_name]} + def getPlamoModelList(self, *args, **kwargs) -> dict: + return {"status":200, "result": config.PLAMO_MODEL_LIST} + + def setPlamoModel(self, data, *args, **kwargs) -> dict: + printLog("Set Plamo Model", data) + try: + data = str(data) + result = model.authenticationTranslatorPlamoAuthKey(auth_key=config.AUTH_KEYS["Plamo_API"], model_name=data) + if result is True: + config.PLAMO_MODEL = data + response = {"status":200, "result":config.PLAMO_MODEL} + else: + response = { + "status":400, + "result":{ + "message":"Plamo model is not valid", + "data": config.PLAMO_MODEL + } + } + except Exception as e: + errorLogging() + response = { + "status":400, + "result":{ + "message":f"Error {e}", + "data": config.PLAMO_MODEL + } + } + return response + def getPlamoAuthKey(self, *args, **kwargs) -> dict: return {"status":200, "result":config.AUTH_KEYS["Plamo_API"]} @@ -1439,14 +1469,24 @@ class Controller: translator_name = "Plamo_API" try: data = str(data) - if len(data) == 32: - key = data - auth_keys = config.AUTH_KEYS - auth_keys[translator_name] = key - config.AUTH_KEYS = auth_keys - config.SELECTABLE_TRANSLATION_ENGINE_STATUS[translator_name] = True - self.updateTranslationEngineAndEngineList() - response = {"status":200, "result":config.AUTH_KEYS[translator_name]} + if len(data) == 72: + result = model.authenticationTranslatorPlamoAuthKey(auth_key=data, model_name=config.PLAMO_MODEL) + if result is True: + key = data + auth_keys = config.AUTH_KEYS + auth_keys[translator_name] = key + config.AUTH_KEYS = auth_keys + config.SELECTABLE_TRANSLATION_ENGINE_STATUS[translator_name] = True + self.updateTranslationEngineAndEngineList() + response = {"status":200, "result":config.AUTH_KEYS[translator_name]} + else: + response = { + "status":400, + "result":{ + "message":"Authentication failure of plamo auth key", + "data": config.AUTH_KEYS[translator_name] + } + } else: response = { "status":400, @@ -2305,7 +2345,7 @@ class Controller: printLog("Start check Plamo API Key") config.SELECTABLE_TRANSLATION_ENGINE_STATUS[engine] = False if config.AUTH_KEYS[engine] is not None: - if model.authenticationTranslatorPlamoAuthKey(auth_key=config.AUTH_KEYS[engine]) is True: + if model.authenticationTranslatorPlamoAuthKey(auth_key=config.AUTH_KEYS[engine], model=config.PLAMO_MODEL) is True: config.SELECTABLE_TRANSLATION_ENGINE_STATUS[engine] = True else: # error update Auth key diff --git a/src-python/mainloop.py b/src-python/mainloop.py index 0791b436..0515a011 100644 --- a/src-python/mainloop.py +++ b/src-python/mainloop.py @@ -168,6 +168,8 @@ mapping = { "/set/data/deepl_auth_key": {"status": False, "variable":controller.setDeeplAuthKey}, "/delete/data/deepl_auth_key": {"status": False, "variable":controller.delDeeplAuthKey}, + "/get/data/plamo_model_list": {"status": False, "variable":controller.getPlamoModelList}, + "/set/data/plamo_model": {"status": False, "variable":controller.setPlamoModel}, "/get/data/plamo_auth_key": {"status": False, "variable":controller.getPlamoAuthKey}, "/set/data/plamo_auth_key": {"status": False, "variable":controller.setPlamoAuthKey}, "/delete/data/plamo_auth_key": {"status": False, "variable":controller.delPlamoAuthKey}, diff --git a/src-python/model.py b/src-python/model.py index 70f828e8..9c2394a1 100644 --- a/src-python/model.py +++ b/src-python/model.py @@ -136,12 +136,12 @@ class Model: del self.keyword_processor self.keyword_processor = KeywordProcessor() - def authenticationTranslatorDeepLAuthKey(self, auth_key): + def authenticationTranslatorDeepLAuthKey(self, auth_key: str) -> bool: result = self.translator.authenticationDeepLAuthKey(auth_key) return result - def authenticationTranslatorPlamoAuthKey(self, auth_key): - result = self.translator.authenticationPlamoAuthKey(auth_key) + def authenticationTranslatorPlamoAuthKey(self, auth_key: str, model: str) -> bool: + result = self.translator.authenticationPlamoAuthKey(auth_key, model=model) return result def startLogger(self): diff --git a/src-python/models/translation/translation_plamo.py b/src-python/models/translation/translation_plamo.py index 0d88e8bc..31046f41 100644 --- a/src-python/models/translation/translation_plamo.py +++ b/src-python/models/translation/translation_plamo.py @@ -1,10 +1,15 @@ from langchain_openai import ChatOpenAI +from pydantic import SecretStr + +_MODELS = [ + "plamo-2.0-prime" + ] class PlamoClient: - def __init__(self, api_key: str): + def __init__(self, api_key: str = "", model: str = "plamo-2.0-prime"): self.api_key = api_key self.base_url = "https://api.platform.preferredai.jp/v1" - self.model = "plamo-2.0-prime" + self.model = model self.supported_languages = """ English Japanese @@ -47,10 +52,55 @@ class PlamoClient: base_url=self.base_url, model=self.model, streaming=True, - openai_api_key=self.api_key, + api_key=SecretStr(self.api_key), ) - def translate_text(self, text: str, input_lang: str, output_lang: str): + def getListModels(self) -> list[str]: + return _MODELS + + def getAuthKey(self) -> str: + """現在のAuthKeyを取得する""" + return self.api_key + + def getModel(self) -> str: + """現在のモデルを取得する""" + return self.model + + def setAuthKey(self, api_key: str) -> bool: + """AuthKeyを設定し、成功したかどうかを返す""" + try: + self.api_key = api_key + self.plamo_llm = ChatOpenAI( + base_url=self.base_url, + model=self.model, + streaming=True, + api_key=SecretStr(self.api_key), + ) + return True + except Exception as e: + print(f"Error setting AuthKey: {e}") + return False + + def setModel(self, model: str) -> bool: + """モデルを設定し、成功したかどうかを返す""" + if model not in _MODELS: + print(f"Model {model} is not in the supported model list.") + return False + + try: + self.model = model + self.plamo_llm = ChatOpenAI( + base_url=self.base_url, + model=self.model, + streaming=True, + api_key=SecretStr(self.api_key), + ) + return True + except Exception as e: + print(f"Error setting model: {e}") + return False + + def translate(self, text: str, input_lang: str, output_lang: str) -> str: messages = [ { "role": "system", @@ -63,18 +113,45 @@ class PlamoClient: output = "" for chunk in self.plamo_llm.stream(messages): - output += chunk.content + if isinstance(chunk.content, str): + output += chunk.content + elif isinstance(chunk.content, list): + for item in chunk.content: + if isinstance(item, str): + output += item + elif isinstance(item, dict): + if "content" in item and isinstance(item["content"], str): + output += item["content"] return output[:-1] + def checkAuthKey(self) -> bool: + try: + self.setModel(self.model) + self.translate("Hello World", input_lang="English", output_lang="Japanese") + return True + except Exception: + return False if __name__ == "__main__": + AUTH_KEY = "AUTH_KEY" text = """ 毎朝コーヒーを入れるのがささやかな楽しみになってる """ input_lang = "Japanese" output_lang = "English" - plamo_client = PlamoClient(api_key="AUTH_KEY") - translated_text = plamo_client.translate_text(text, input_lang, output_lang) - print(translated_text) \ No newline at end of file + plamo_client = PlamoClient(api_key=AUTH_KEY) + + print("model list:", plamo_client.getListModels()) + print("AuthKey:", plamo_client.getAuthKey()) + print("Model:", plamo_client.getModel()) + print(f"set model: {plamo_client.setModel('plamo-2.0-prime')}") + print(f"set AuthKey: {plamo_client.setAuthKey(AUTH_KEY)}") + print(f"check AuthKey: {plamo_client.checkAuthKey()}") + + try: + translated_text = plamo_client.translate(text, input_lang, output_lang) + print(translated_text) + except Exception: + print("Invalid API key. Please check your credentials.") \ No newline at end of file diff --git a/src-python/models/translation/translation_translator.py b/src-python/models/translation/translation_translator.py index 212dc56c..e187d52a 100644 --- a/src-python/models/translation/translation_translator.py +++ b/src-python/models/translation/translation_translator.py @@ -35,10 +35,10 @@ class Translator(): self.is_loaded_ctranslate2_model = False self.is_enable_translators = ENABLE_TRANSLATORS - def authenticationDeepLAuthKey(self, authkey): + def authenticationDeepLAuthKey(self, auth_key: str) -> bool: result = True try: - self.deepl_client = DeepLClient(authkey) + self.deepl_client = DeepLClient(auth_key) self.deepl_client.translate_text("Hello World", target_lang="EN-US") except Exception: errorLogging() @@ -46,11 +46,11 @@ class Translator(): result = False return result - def authenticationPlamoAuthKey(self, authkey): + def authenticationPlamoAuthKey(self, auth_key: str, model: str) -> bool: result = True try: - self.plamo_client = PlamoClient(authkey) - self.plamo_client.translate_text("Hello World", target_lang="English") + self.plamo_client = PlamoClient(auth_key, model=model) + self.plamo_client.checkAuthKey() except Exception: errorLogging() self.plamo_client = None @@ -153,7 +153,7 @@ class Translator(): if self.plamo_client is None: result = False else: - result = self.plamo_client.translate_text( + result = self.plamo_client.translate( message, input_lang=source_language, output_lang=target_language,