Spaces:
Running
Running
| 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() |