from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

tokenizer = AutoTokenizer.from_pretrained("merve/chatgpt-prompts-bart-long")
model = AutoModelForSeq2SeqLM.from_pretrained("merve/chatgpt-prompts-bart-long", from_tf=True)

def generate(prompt):
    profession = "example profession"  # Replace with your desired profession
    description = "example"  # Replace with your desired description

    formatted_prompt = f"You are an expert at {profession} who always {description}.\n\n"

    batch = tokenizer(formatted_prompt, return_tensors="pt")
    generated_ids = model.generate(batch["input_ids"], max_length=200, num_return_sequences=1)
    output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
    return output[0]

input_component = gr.Textbox(label="Input a persona, e.g. photographer", value="photographer")
output_component = gr.Textbox(label="Prompt")
examples = [["photographer"], ["developer"]]
description = "This app generates ChatGPT prompts. Enter a persona you want the prompt to be based on."
gr.Interface(generate, inputs=input_component, outputs=output_component, examples=examples, title="👨🏻‍🎤 ChatGPT Prompt Generator 👨🏻‍🎤", description=description).launch()