history_mistery / app.py
SaviAnna's picture
Update app.py
49e7ba4
raw
history blame
1.74 kB
import transformers
import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import numpy as np
from PIL import Image
# import plotly.figure_factory as ff
st.write("""
# History Mistery
""")
# image = Image.open('data-scins.jpeg')
# st.image(image, caption='Current mood')
# Загрузка модели и токенизатора
model = GPT2LMHeadModel.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
# Функция для генерации текста
def generate_text(prompt):
# Преобразование входной строки в токены
input_ids = tokenizer.encode(prompt, return_tensors='pt')
# Генерация текста
output = model.generate(input_ids=input_ids, max_length=70, num_beams=5, do_sample=True,
temperature=1.0, top_k=50, top_p=0.6, no_repeat_ngram_size=3,
num_return_sequences=3)
# Декодирование сгенерированного текста
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
return generated_text
# Streamlit приложение
def main():
st.title("GPT-3 Text Generation")
# Ввод строки пользователем
prompt = st.text_area("Enter your prompt:", value="В средние века на руси")
# Генерация текста по введенной строке
generated_text = generate_text(prompt)
# Отображение сгенерированного текста
st.subheader("Generated Text:")
st.write(generated_text)
if name == "__main__":
main()