import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def run():
    @st.cache_resource
    def load_model():
        model = AutoModelForCausalLM.from_pretrained("gpt")
        tokenizer = AutoTokenizer.from_pretrained("gpt")
        return model, tokenizer

    model, tokenizer = load_model()
    model.eval()

    # === Интерфейс ===
    st.title("📝 Генерация текста GPT-моделью")
    prompt = st.text_area("Ввод пользователя (prompt)", "Не думай, друг, что время лечит —")

    max_length = st.slider("Максимальная длина генерации", 20, 300, 100)
    num_return_sequences = st.slider("Количество генераций", 1, 5, 1)
    temperature = st.slider("Temperature", 0.1, 2.0, 1.0)
    top_k = st.slider("Top-k", 0, 100, 50)
    top_p = st.slider("Top-p", 0.0, 1.0, 0.95)

    generate_btn = st.button("Сгенерировать")

    if generate_btn:
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=max_length,
                num_return_sequences=num_return_sequences,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        st.subheader("📄 Сгенерированный текст:")
        for i, output in enumerate(outputs):
            text = tokenizer.decode(output, skip_special_tokens=True)
            st.markdown(f"**Вариант {i+1}:**\n\n{text}")