"""Minimal example. Usage: python run_example.py "Hello, how are you?" sv"""
import sys, os, torch, sentencepiece as spm
from modeling_flash import NordicFlash
H = os.path.dirname(__file__)
LANG = {"en":65000,"sv":65001,"da":65002,"nb":65003,"nn":65004,"fi":65005,"is":65006}
text = sys.argv[1] if len(sys.argv) > 1 else "Hello, how are you today?"
tgt  = sys.argv[2] if len(sys.argv) > 2 else "sv"
sp = spm.SentencePieceProcessor(); sp.load(H + "/nordic_unigram_65k.model")
m = NordicFlash.from_checkpoint(H + "/model.safetensors", device="cuda" if torch.cuda.is_available() else "cpu")
out = m.translate(sp.encode(text, out_type=int), LANG[tgt])
print(f"[{tgt}] {sp.decode(out)}")