Spaces:
Sleeping
Sleeping
File size: 1,504 Bytes
8fb2bb2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import streamlit as st
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import textwrap
st.title('GPT2 trained on tg chat')
model_directory = 'finetuned/' # Directory where the model is located
model = GPT2LMHeadModel.from_pretrained(model_directory, use_safetensors=True)
tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
def predict(text, max_len=100, num_beams=10, temperature=1.5, top_p=0.7):
with torch.inference_mode():
prompt = text
prompt = tokenizer.encode(prompt, return_tensors='pt')
out = model.generate(
input_ids=prompt,
max_length=max_len,
num_beams=num_beams,
do_sample=True,
temperature=temperature,
top_p=top_p,
no_repeat_ngram_size=1,
num_return_sequences=1,
).cpu().numpy()
return textwrap.fill(tokenizer.decode(out[0]))
prompt = st.text_input("Твоя фраза")
col = st.columns(4)
with col[0]:
max_len = st.slider("Text len", 20, 200, 100)
with col[1]:
num_beams = st.slider("Beams", 0.1, 1., 0.5)
with col[2]:
temperature = st.slider("Temperature", 0.1, 0.9, 0.35)
with col[3]:
top_p = st.slider("Top-p", 0.1, 1.0, 0.7)
submit = st.button('Сгенерировать ответ')
if submit:
if prompt:
pred = predict(prompt, max_len=max_len, num_beams=int(num_beams * 20), temperature=(1-temperature) * 5, top_p=top_p)
st.write(pred)
|