import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline
import pickle

tokenizer = AutoTokenizer.from_pretrained("daspartho/subreddit-predictor")
model = AutoModelForSequenceClassification.from_pretrained("daspartho/subreddit-predictor") # i've uploaded the model on HuggingFace :)

with open('labels.bin', 'rb') as f:
    label_map = pickle.load(f)

pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer, top_k=3)

def classify_text(plot):
    predictions = pipe(plot)[0]
    return {label_map[pred['label']]: float(pred['score']) for pred in predictions}

examples = [
            ["My frying pan dried with a spot that looks like a pig"], 
            ["Adult peer pressure is hearing your neighbor mowing so you decide you better mow too"],
            ['a bear walks into a bar and says, "give me a whiskey and... cola"'],
            ["Worst Celebrity Private Jet CO2 Emissions Offenders (2022)"],
            ["Billionaire No More: Patagonia Founder Gives Away the Company - Profits will now go towards climate action"],
            ["How an AI imagines infinity"],
            ["Understanding consciousness is more important than ever at the dawn of this AI age"],
            ["White T-Shirt Inspiration Album"],
            ["Be proud of your progress."],
            ["Avocado toast with poached egg and hot sauce."],
            ]

iface = gr.Interface(
    description = "Enter a title for a reddit post, and the model will attempt to predict the subreddit.",
    article = "<p style='text-align: center'><a href='https://github.com/daspartho/predict-subreddit' target='_blank'>Github</a></p>",
    fn=classify_text, 
    inputs=gr.inputs.Textbox(label="Type the title here"), 
    outputs=gr.outputs.Label(label='What the model thinks'),
    examples=examples
    )
iface.launch()