Spaces:
No application file
No application file
| 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 | |