from llama_cpp import Llama
from huggingface_hub import hf_hub_download



# tencent/HY-MT1.5-1.8B
# 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,
):
    """
    Translate a short text segment to a target language using a Llama model.

    This function builds a simple prompt and sends it to the module-level `llm`
    instance to obtain a direct translation of the provided text. The returned
    value is the translated text string, without additional explanation.

    Parameters:
    - text (str): Text to be translated — a sentence or short paragraph.
    - target_language (str): Target language (e.g., "Portuguese", "Spanish").
    - temperature (float): Controls the randomness of the output; lower values
        make the translation more deterministic.

    Returns:
    - str: Translated text (trimmed with `strip()`), ready for use.

    Notes and considerations:
    - The `stop=["\n"]` parameter is set to attempt to stop at the first
        newline; for multi-line translations consider adjusting or removing `stop`.

    Example:
    translated = run("It's on the house.", "Portuguese", temperature=0.2)
    """
    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