Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
2
3# Load model and tokenizer
4tokenizer = AutoTokenizer.from_pretrained("Lahad/nllb200-francais-wolof")
5model = AutoModelForSeq2SeqLM.from_pretrained("Lahad/nllb200-francais-wolof")
6
7# Translation function
8def translate(text, max_length=128):
9 inputs = tokenizer(
10 text,
11 max_length=max_length,
12 padding="max_length",
13 truncation=True,
14 return_tensors="pt"
15 )
16
17 outputs = model.generate(
18 input_ids=inputs["input_ids"],
19 attention_mask=inputs["attention_mask"],
20 forced_bos_token_id=tokenizer.convert_tokens_to_ids("wol_Latn"),
21 max_length=max_length
22 )
23
24 return tokenizer.decode(outputs[0], skip_special_tokens=True)