import gradio as gr
import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed


set_seed(42)
tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", torch_dtype=torch.bfloat16)


def Bemenet(bemenet):
    prompt = "<human>: Who is Alan Turing?\n<bot>:"
    inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
    input_length = inputs.input_ids.shape[1]
    outputs = model.generate(
        **inputs, max_new_tokens=128, do_sample=True, temperature=0.7, top_p=0.7, top_k=50, return_dict_in_generate=True
    )
    token = outputs.sequences[0, input_length:]
    output_str = tokenizer.decode(token)
    return output_str


interface = gr.Interface(fn=Bemenet,
                        title="Cím..",
                        description="Leírás..",
                        inputs="text",
                        outputs="text")

interface.launch()