byt5-khmer-en-alt

import torch
import unicodedata
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 1. Load Model
MODEL_ID = "Darayut/byt5-khmer-en-alt"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID).to(device)

def translate(text):
    # --- PREPROCESSING (Must match training) ---
    # 1. Normalize to NFC (Fixes hidden Khmer vowel issues)
    # This is crucial for ByT5 as it reads raw bytes.
    text = unicodedata.normalize("NFC", text.strip())
    
    # 2. Tokenize
    inputs = tokenizer(text, return_tensors="pt").input_ids.to(device)
    
    # 3. Generate
    # max_length=128 is usually enough for English sentences
    outputs = model.generate(inputs, max_length=384)
    
    # 4. Decode
    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translation

# --- Example Usage ---
khmer_text = "រដ្ឋាភិបាលប្រទេសអង់គ្លេសបានកំណត់ទិសដៅអតិផរណាចំនួន2% ដែលបានកើនលើសជាង1%។"
result = translate(khmer_text)
print(f"Output: {result}")