import streamlit as st
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer


model_name_or_path = "sberbank-ai/rugpt3small_based_on_gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path)
model =  GPT2LMHeadModel.from_pretrained(
    model_name_or_path,
    output_attentions = False,
    output_hidden_states = False,
)

# Загрузка сохраненных весов
model_weights_path = "nlp_project/hunter_generator.pt"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load(model_weights_path, map_location=device))
model.eval()
def generate_text(user_input, model=model, tokenizer=tokenizer):
    input_ids = tokenizer.encode(user_input, return_tensors="pt")
    with torch.no_grad():
        out = model.generate(
            input_ids,
            do_sample=True,
            num_beams=3,
            temperature=1.05,
            top_p=.8,
            max_length=50,
        )
    generated_text = list(map(tokenizer.decode, out))[0]
    return generated_text