File size: 1,349 Bytes
1d6cf0f |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 |
import gradio as gr
from transformers import pipeline
# Load the zero-shot classification pipeline
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
def classify_text(line_item, classes):
# Split the classes string into a list
class_list = classes.split(',')
# Perform classification
results = classifier(line_item, class_list, multi_class=True)
# Prepare the output as a dictionary {label: score}
output = {label: round(score, 4) for label, score in zip(results['labels'], results['scores'])}
return output
# Define Gradio interface with example
interface = gr.Interface(
classify_text,
[
gr.Textbox(lines=2, placeholder="Enter Line Item Here...", label="Line Item"),
gr.Textbox(placeholder="Enter Classes Here, Separated by Commas", label="Classes")
],
gr.Label(num_top_classes=None, label="Class Probability Scores"),
title="Bad Stuff, But So Good.",
description="A zero-shot classification app using facebook/bart-large-mnli model to classify text into given categories.",
examples=[
["wijn glas x3 $18", "tobacco,alcohol"], # Example input
["Stoofvlees $25", "tobacco,alcohol"],
["Marlboro 10$", "tobacco,alcohol"],
]
)
# Launch the application
if __name__ == "__main__":
interface.launch()
|