| """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)}") | |