| import sentencepiece as spm |
| from vllm import LLM, SamplingParams |
| from vllm.inputs import TokensPrompt |
|
|
| SPM = "/model/nordic_unigram_65k.model" |
| LANG = {"en": 65000, "sv": 65001, "da": 65002, "nb": 65003, |
| "nn": 65004, "fi": 65005, "is": 65006} |
| BOS, EOS, EOS_SRC = 1, 2, 65007 |
|
|
| def main(): |
| sp = spm.SentencePieceProcessor(); sp.load(SPM) |
| llm = LLM(model="/model", trust_remote_code=True, skip_tokenizer_init=True, |
| dtype="bfloat16", max_model_len=512, gpu_memory_utilization=0.55, |
| enforce_eager=True) |
|
|
| def translate(text, tgt): |
| ids = [BOS, LANG[tgt]] + sp.encode(text, out_type=int) + [EOS_SRC] |
| sp_out = llm.generate(TokensPrompt(prompt_token_ids=ids), |
| SamplingParams(temperature=0.0, max_tokens=64, |
| stop_token_ids=[EOS])) |
| toks = list(sp_out[0].outputs[0].token_ids) |
| return sp.decode([t for t in toks if t < 65000]) |
|
|
| print("\n===== vLLM TRANSLATIONS =====") |
| for txt, tgt in [("Hello, how are you today?", "sv"), |
| ("The weather is nice and the sun is shining.", "sv"), |
| ("I would like to order a coffee, please.", "nb"), |
| ("Thank you very much for your help.", "is")]: |
| print(f"[en->{tgt}] {txt}\n -> {translate(txt, tgt)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
