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()