import gradio as gr

from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch


def vipe_generate(text, model, tokenizer,device,do_sample,top_k=100, epsilon_cutoff=.00005, temperature=1):
    #mark the text with special tokens
    text=[tokenizer.eos_token +  i + tokenizer.eos_token for i in text]
    batch=tokenizer(text, padding=True, return_tensors="pt")

    input_ids = batch["input_ids"].to(device)
    attention_mask = batch["attention_mask"].to(device)

    #how many new tokens to generate at max
    max_prompt_length=50

    generated_ids = model.generate(input_ids=input_ids,attention_mask=attention_mask, max_new_tokens=max_prompt_length, do_sample=do_sample,top_k=top_k, epsilon_cutoff=epsilon_cutoff, temperature=temperature)
    #return only the generated prompts
    pred_caps = tokenizer.batch_decode(generated_ids[:, -(generated_ids.shape[1] - input_ids.shape[1]):], skip_special_tokens=True)

    return pred_caps[0]

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = GPT2LMHeadModel.from_pretrained('fittar/ViPE-M-CTX7')
model.to(device)
tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
tokenizer.pad_token = tokenizer.eos_token

def generate(text):
    result =vipe_generate([text],model,tokenizer,do_sample=True,device=device)
    return result


examples = [
    ["brave, fantasy"],
    ["She felt like a flower in December"],
    ["2+2=4? hmm.."]
]

title = "ViPE: Visualize Pretty-much Everything"
description = 'ViPE is the first automated model for translating any arbitrary piece of text into a visualizable prompt. It helps any text-to-image model in figurative or non-lexical language visualizations. To learn more about the model, [click here](https://huggingface.co/fittar/ViPE-M-CTX7).<br>'
txt = gr.Textbox(lines=1, label="Arbitrary Input Text", placeholder="Initial text")
out = gr.Textbox(lines=4, label="Generated Prompt for Visualizations")

demo = gr.Interface(
    fn =generate,
    inputs=txt,
    outputs=out,
    examples=examples,
    title=title,
    description=description,
    theme="default",
    cache_examples="never"
)

demo.launch(enable_queue=True, debug=True)