Views
No views yet
pip install -U bitsandbytes1import torch
2import transformers
3
4# Load the 4-bit quantized model and tokenizer
5model_4bit = transformers.M2M100ForConditionalGeneration.from_pretrained(
6 "Sunbird/translate-nllb-1.3b-salt-4bit",
7 device_map="auto"
8)
9tokenizer = transformers.NllbTokenizer.from_pretrained("Sunbird/translate-nllb-1.3b-salt")
10
11# Define the text and language parameters
12text = 'Where is the hospital?'
13source_language = 'eng'
14target_language = 'lug'
15
16# Mapping for language tokens
17language_tokens = {
18 'eng': 256047,
19 'ach': 256111,
20 'lgg': 256008,
21 'lug': 256110,
22 'nyn': 256002,
23 'teo': 256006,
24}
25
26# Prepare device and tokenize the input text
27device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
28inputs = tokenizer(text, return_tensors="pt").to(device)
29inputs['input_ids'][0][0] = language_tokens[source_language]
30
31# Generate the translation with beam search
32translated_tokens = model_4bit.to(device).generate(
33 **inputs,
34 forced_bos_token_id=language_tokens[target_language],
35 max_length=100,
36 num_beams=5,
37)
38
39# Decode and print the translated result
40result = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
41print(result)
42# Expected output: "Eddwaliro liri ludda wa?"