from logging import getLogger
from llama_cpp import Llama
from functools import lru_cache

logger = getLogger(__name__)

class QwenTranslator:
    def __init__(self, model_path, system_prompt_en="", system_prompt_zh="") -> None:
        self.llm = Llama(
        model_path=model_path,
        chat_format="chatml",
        verbose=False)
        self.sys_prompt_en = system_prompt_en
        self.sys_prompt_zh = system_prompt_zh

    def to_message(self, prompt, src_lang, dst_lang):
        """构造提示词"""
        return [
            {"role": "system", "content": self.sys_prompt_en if src_lang == "en" else self.sys_prompt_zh},
            {"role": "user", "content": prompt},
        ]

    @lru_cache(maxsize=10)
    def translate(self, prompt, src_lang, dst_lang) -> str:
        message = self.to_message(prompt, src_lang, dst_lang)
        output = self.llm.create_chat_completion(messages=message, temperature=0)
        return output['choices'][0]['message']['content']

    def __call__(self, prompt,*args, **kwargs):
        return self.llm(
            prompt,
            *args,
            **kwargs
        )