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}

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'),
    )
iface.launch()