epo_lstm / example.py
timcryt's picture
Upload 2 files
eff5017 verified
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'))