import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "oddadmix/Emhotob-10M-Darija-English-v2"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16).to("cuda").eval()
SYSTEM = "You are a professional translator. Translate the Moroccan Darija text into English."
def translate(text, system=SYSTEM):
prompt = (f"<|im_start|>system\n{system}<|im_end|>\n"
f"<|im_start|>user\n{text.strip()}<|im_end|>\n<|im_start|>assistant\n")
ids = tok(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
if tok.bos_token_id is not None:
bos = torch.tensor([[tok.bos_token_id]], device=model.device)
ids["input_ids"] = torch.cat([bos, ids["input_ids"]], dim=1)
ids["attention_mask"] = torch.cat([torch.ones_like(bos), ids["attention_mask"]], dim=1)
out = model.generate(**ids, max_new_tokens=256, do_sample=False,
eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id)
return tok.decode(out[0, ids["input_ids"].size(1):], skip_special_tokens=True).strip()