import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
LANG_NAMES = ["English", "German"]
MODEL_ID = "finnianx/Gros-Michel-Translator"
TOKENIZER_ID = "finnianx/Gros-Michel-90m-Base-v2"
DEVICE = "cuda"

tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32).to(DEVICE)

def translate(source: str, target: str, text: str) -> str:
    prompt = f"Translate {source} to {target}:\n{source}: {text}\n{target}:"
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    inputs.pop("token_type_ids", None)
    out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
    return tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()