new_space_nlp / pages /seduce.py
Teery's picture
first try
f1f9e9e
raw
history blame contribute delete
No virus
1.3 kB
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])