student_assistant / app_test.py
prouditalian's picture
Create app_test.py
8b05fe9
raw
history blame contribute delete
637 Bytes
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import streamlit as st
tokenizer = GPT2Tokenizer.from_pretrained("sberbank-ai/mGPT")
model = GPT2LMHeadModel.from_pretrained("sberbank-ai/mGPT")
text = "Александр Сергеевич Пушкин родился в "
input_ids = tokenizer.encode(text, return_tensors="pt").cuda(device)
out = model.generate(
input_ids,
min_length=100,
max_length=100,
eos_token_id=5,
pad_token=1,
top_k=10,
top_p=0.0,
no_repeat_ngram_size=5
)
generated_text = list(map(tokenizer.decode, out))[0]
st.write(generated_text)