new_space_nlp / pages /seduce.py
Teery's picture
first try
f1f9e9e
import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments
import torch
st.header('Генератор плохих подкатов')
tokenizer = GPT2Tokenizer.from_pretrained('SeducerWeight/repo_name')
model = GPT2LMHeadModel.from_pretrained('SeducerWeight/repo_name', from_tf=False)
text = st.text_input("Кринжовое начало")
max_length = st.slider('Выберите длину', 30, 150, 40, 5)
temp = st.slider('Выберите температуру', 1.5, 4.0, 2.0, 0.25)
top_p = st.slider('Выберите top_p', 1.5, 4.0, 2.0, 0.25)
if text and max_length and temp and top_p:
input_ids = tokenizer.encode(text, return_tensors="pt")
model.eval()
with torch.no_grad():
out = model.generate(input_ids,
do_sample=True,
num_beams=3,
temperature=temp,
top_p=top_p,
max_length=max_length,
)
generated_text = list(map(tokenizer.decode, out))
count = st.number_input('Введите количество подкатов', 1, 5, 1)
st.write(generated_text[0].split('\n')[count-1])