Spaces:

drift-ai
/

bad-stuff

Sleeping

bad-stuff / app.py

Vincent Claes

inital commit

1d6cf0f 3 months ago

No virus

1.35 kB

	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()