File size: 780 Bytes
4cf4a9e ab6f576 4cf4a9e ab6f576 4cf4a9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
from transformers import pipeline
import gradio as gr
model_checkpoint = "MuntasirHossain/distilbert-finetuned-ag-news"
model = pipeline("text-classification", model=model_checkpoint)
def predict(prompt):
completion = model(prompt)[0]["label"]
return completion
description = "This AI model is trained to classify news articles into four categories: World, Sports, Business and Science/Tech."
title = "Classify Your Articles"
theme = "peach"
examples=[["Global Retail Giants Gear Up for Record-Breaking Holiday Sales Season Amidst Supply Chain Challenges and Rising Consumer Demand."], ["Business"]]
gr.Interface(fn=predict,
inputs="textbox",
outputs="text",
title=title,
theme = theme,
description=description,
examples=examples,
).launch() |