from llama_cpp import Llama
from huggingface_hub import hf_hub_download



# tencent/HY-MT1.5-1.8B-GGUF, HY-MT1.5-1.8B-Q6_K.gguf, HY-MT1.5-1.8B-Q4_K_M.gguf
# mradermacher/HY-MT1.5-1.8B-GGUF, HY-MT1.5-1.8B.Q5_K_S.gguf
repo_id = "tencent/HY-MT1.5-1.8B-GGUF"
filename = "HY-MT1.5-1.8B-Q4_K_M.gguf"
local_dir = "./model"
model_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)

llm = Llama(
    model_path=model_path,
    n_ctx=4096,      # Corresponde ao -n 4096
    verbose=False    # Reduz o log excessivo no console
)



def run(
    text = "It’s on the house.",
    target_language = "Portuguese",
    temperature = 0.3,
):
    prompt = f"Translate the following segment into {target_language}, without additional explanation.\n\n{text}"

    output = llm(
        prompt,
        max_tokens=4096,
        temperature=temperature,      # --temp 0.7
        top_k=20,             # --top-k 20
        top_p=0.6,             # --top-p 0.6
        repeat_penalty=1.05,  # --repeat-penalty 1.05
        stop=["\n"]           # Opcional: para parar após a tradução
    )

    result = output["choices"][0]["text"].strip()
    return result