app.py

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

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 vipe_generate(text):
    result =generate(text,model,tokenizer,do_sample=True,device=device)
    return result

examples = [
    ["Is string theory right?"],
    ["She felt like a flower in December"],
    ["2+2=4?"],
]

demo = gr.Interface(
    fn=vipe_generate,
    inputs=gr.inputs.Textbox(lines=5, label="Arbitrary Input Text"),
    outputs=gr.outputs.Textbox(label="Generated Prompt for Visualizations"),
    examples=examples
)

demo.launch()