import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import random
device = 'cpu'

def ans(question ):
    description=''
    category=''
    seed = random.randint(1, 10000000)
    print(f'Seed: {seed}')
    torch.manual_seed(seed)
    
    inp = tokenizer.encode(f'Вопрос: {question}\nОписание: {description}\nОтвет:',return_tensors="pt").to(device)
    print('question',question)
    gen = model.generate(inp, do_sample=True, top_p=0.9, temperature=0.86, max_new_tokens=100, repetition_penalty=1.2) #, stop_token="<eos>")
    
    gen = tokenizer.decode(gen[0])
    gen = gen[:gen.index('<eos>') if '<eos>' in gen else len(gen)]
    gen = gen.split('Ответ:')[1]
    return gen







# Download checkpoint:
checkpoint = "its5Q/rugpt3large_mailqa"
tokenizer =  AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint)
model = model.eval()

# Gradio

title = "Ответы на главные вопросы жизни, вселенной и вообще"
description = "ruGPT large дообученная на датасете https://www.kaggle.com/datasets/atleast6characterss/otvetmailru-solved-questions "
article = "<p style='text-align: center'><a href='https://github.com/NeuralPushkin/MailRu_Q-A'>Github with fine-tuning ruGPT3large on QA</a></p> Cозданно при поддержке <p style='text-align: center'><a href='https://t.me/lovedeathtransformers'>Love Death Transformers</a></p>"
examples = [
            ["Как какать?"]
]

iface = gr.Interface(fn=ans, title=title, description=description, article=article, examples=examples, inputs="text", outputs="text")

if __name__ == "__main__":
    iface.launch()