Views
No views yet
google/byt5-small1import torch
2import unicodedata
3from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
4
5# 1. Load Model
6MODEL_ID = "Darayut/byt5-small-khm-en-translation"
7device = "cuda" if torch.cuda.is_available() else "cpu"
8
9print(f"Loading {MODEL_ID}...")
10tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
11model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID).to(device)
12
13def translate(text):
14 # --- PREPROCESSING (Must match training) ---
15 # 1. Normalize to NFC (Fixes hidden Khmer vowel issues)
16 # This is crucial for ByT5 as it reads raw bytes.
17 text = unicodedata.normalize("NFC", text.strip())
18
19 # 2. Tokenize
20 inputs = tokenizer(text, return_tensors="pt").input_ids.to(device)
21
22 # 3. Generate
23 # max_length=128 is usually enough for English sentences
24 outputs = model.generate(inputs, max_length=128)
25
26 # 4. Decode
27 translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
28 return translation
29
30# --- Example Usage ---
31khmer_text = "ជីវិតគឺជាការធ្វើដំណើរដែលពោរពេញដោយបទពិសោធន៍"
32result = translate(khmer_text)
33print(f"Output: {result}")
34
35# Expected Output: "Life is a journey fulfilled by the experience"unicodedata.normalize("NFC", text).