import gradio as gr
import random
from datasets import load_dataset
from transformers import pipeline

# Load dataset (optional if you want random roast sampling)
dataset = load_dataset("gus-gustavo/reddit_roastme")

# Use DialoGPT from Hugging Face
generator = pipeline("text-generation", model="microsoft/DialoGPT-large")

def roast_me(prompt):
    # Generate roast using DialoGPT
    response = generator(prompt, max_length=100, do_sample=True)[0]["generated_text"]
    return response

# Option B: Just sample from dataset
def random_roast(_):
    return random.choice(dataset["train"])["body"]

# Gradio interface
interface = gr.Interface(
    fn=roast_me,   # swap to random_roast if you want dataset-only
    inputs="text",
    outputs="text",
    title="RoastBot",
    description="A chatbot that roasts you using Microsoft DialoGPT or random roasts from Reddit's RoastMe dataset."
)

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