from epo_lstm import LSTMGeneratorPipeline, EOLSTMGenerator

from transformers import PreTrainedTokenizerFast

generator = LSTMGeneratorPipeline(
    model=EOLSTMGenerator.from_pretrained("./version_1"),
    tokenizer=PreTrainedTokenizerFast(
        tokenizer_file="./version_1/tokenizer.json",
        bos_token="[BOS]",
        eos_token="\n",
        unk_token="[UNK]",
        pad_token="[PAD]",
    ),
)

print(generator('saluton, mi estas'))