from pathlib import Path
import numpy as np
from lib.utils import Timer
from s2ts import S2TS
from s2ts import TaskExecInfo as CTaskExecInfo

MODEL_DIR = Path(r"D:\yujuan\yoyo-translator-win\models\llm\Qwen3-1.7B-int8-ov")

class QwenOv:
    def __init__(self, model_dir=MODEL_DIR):
        with Timer("load LLM"):
            self.instance = S2TS()
            ret = self.instance.start_translate_genai(str(model_dir))
            print(f"model load {'success' if ret else 'failed'}")
        self._warm_up()
    def _warm_up(self):
        self.translate("How are you?", "en", "zh")

    def translate(self, prompt, src_lang, dst_lang):
        task_info = CTaskExecInfo()
        task_info.transcribe_content = prompt
        task_info.audio_language=src_lang
        task_info.translate_language=dst_lang

        with Timer("LLM inference") as t:
            self.instance.put_llm(task_info)
            res: CTaskExecInfo = self.instance.get_llm(0)
        return res.translate_content, t.duration