Spaces:

alexandrainst
/

zero-shot-classification

Running

App Files Files Community

zero-shot-classification / app.py

saattrupdan

feat: Update deps, sort out deprecation warnings

af97119 about 1 year ago

raw history blame contribute delete

No virus

10.1 kB

	"""Gradio app that showcases Scandinavian zero-shot text classification models."""

	from typing import Dict, Tuple
	import gradio as gr
	from gradio.components import Dropdown, Textbox, Row, Column, Button, Label, Markdown
	from types import MethodType
	from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
	from luga import language as detect_language
	import torch
	import re
	import os
	import torch._dynamo


	def main():
	# Disable tokenizers parallelism
	os.environ["TOKENIZERS_PARALLELISM"] = "false"

	# Load the zero-shot classification pipeline
	global classifier, model, tokenizer
	model_id = "alexandrainst/scandi-nli-large"
	model = AutoModelForSequenceClassification.from_pretrained(model_id)
	tokenizer = AutoTokenizer.from_pretrained(model_id)
	model = torch.compile(model=model, backend="aot_eager")
	model.eval()
	classifier = pipeline("zero-shot-classification", model=model, tokenizer=tokenizer)
	classifier.get_inference_context = MethodType(
	lambda self: torch.no_grad, classifier
	)

	# Create dictionary of descriptions for each task, containing the hypothesis template
	# and candidate labels
	task_configs: Dict[str, Tuple[str, str, str, str, str, str]] = {
	"Sentiment classification": (
	"Dette eksempel er {}.",
	"positivt, negativt, neutralt",
	"Detta exempel är {}.",
	"positivt, negativt, neutralt",
	"Dette eksemplet er {}.",
	"positivt, negativt, nøytralt",
	),
	"News topic classification": (
	"Denne nyhedsartikel handler primært om {}.",
	"krig, politik, uddannelse, sundhed, økonomi, mode, sport",
	"Den här nyhetsartikeln handlar främst om {}.",
	"krig, politik, utbildning, hälsa, ekonomi, mode, sport",
	"Denne nyhetsartikkelen handler først og fremst om {}.",
	"krig, politikk, utdanning, helse, økonomi, mote, sport",
	),
	"Spam detection": (
	"Denne e-mail ligner {}.",
	"en spam e-mail, ikke en spam e-mail",
	"Det här e-postmeddelandet ser {}.",
	"ut som ett skräppostmeddelande, inte ut som ett skräppostmeddelande",
	"Denne e-posten ser {}.",
	"ut som en spam-e-post, ikke ut som en spam-e-post",
	),
	"Product feedback detection": (
	"Denne kommentar er {}.",
	"en anmeldelse af et produkt, ikke en anmeldelse af et produkt",
	"Den här kommentaren är {}.",
	"en recension av en produkt, inte en recension av en produkt",
	"Denne kommentaren er {}.",
	"en anmeldelse av et produkt, ikke en anmeldelse av et produkt",
	),
	"Define your own task!": (
	"Dette eksempel er {}.",
	"",
	"Detta exempel är {}.",
	"",
	"Dette eksemplet er {}.",
	"",
	),
	}

	def set_task_setup(task: str) -> Tuple[str, str, str, str, str, str]:
	return task_configs[task]

	with gr.Blocks() as demo:

	# Create title and description
	Markdown("# Scandinavian Zero-shot Text Classification")
	Markdown("""
	Classify text in Danish, Swedish or Norwegian into categories, without
	finetuning on any training data!

	Select one of the tasks from the dropdown menu on the left, and try
	entering some input text (in Danish, Swedish or Norwegian) in the input
	text box and press submit, to see the model in action! The labels are
	generated by putting in each candidate label into the hypothesis template,
	and then running the classifier on each label separately. Feel free to
	change the "hypothesis template" and "candidate labels" on the left as you
	please as well, and try to come up with your own tasks too 😊

	_Also, be patient, as this demo is running on a CPU!_
	""")

	with Row():

	# Input column
	with Column():

	# Create a dropdown menu for the task
	dropdown = Dropdown(
	label="Task",
	choices=[
	"Sentiment classification",
	"News topic classification",
	"Spam detection",
	"Product feedback detection",
	"Define your own task!",
	],
	value="Sentiment classification",
	)

	with Row(variant="compact"):
	da_hypothesis_template = Textbox(
	label="Danish hypothesis template",
	value="Dette eksempel er {}.",
	)
	da_candidate_labels = Textbox(
	label="Danish candidate labels (comma separated)",
	value="positivt, negativt, neutralt",
	)

	with Row(variant="compact"):
	sv_hypothesis_template = Textbox(
	label="Swedish hypothesis template",
	value="Detta exempel är {}.",
	)
	sv_candidate_labels = Textbox(
	label="Swedish candidate labels (comma separated)",
	value="positivt, negativt, neutralt",
	)

	with Row(variant="compact"):
	no_hypothesis_template = Textbox(
	label="Norwegian hypothesis template",
	value="Dette eksemplet er {}.",
	)
	no_candidate_labels = Textbox(
	label="Norwegian candidate labels (comma separated)",
	value="positivt, negativt, nøytralt",
	)

	# When a new task is chosen, update the description
	dropdown.change(
	fn=set_task_setup,
	inputs=dropdown,
	outputs=[
	da_hypothesis_template,
	da_candidate_labels,
	sv_hypothesis_template,
	sv_candidate_labels,
	no_hypothesis_template,
	no_candidate_labels,
	],
	)

	# Output column
	with Column():

	# Create a text box for the input text
	input_textbox = Textbox(
	label="Input text", value="Jeg er helt vild med fodbolden 😊"
	)

	with Row():
	clear_btn = Button(value="Clear")
	submit_btn = Button(value="Submit", variant="primary")

	# When the clear button is clicked, clear the input text box
	clear_btn.click(
	fn=lambda _: "", inputs=input_textbox, outputs=input_textbox
	)


	with Column():

	# Create output text box
	output_textbox = Label(label="Result")

	# When the submit button is clicked, run the classifier on the input text
	# and display the result in the output text box
	submit_btn.click(
	fn=classification,
	inputs=[
	input_textbox,
	da_hypothesis_template,
	da_candidate_labels,
	sv_hypothesis_template,
	sv_candidate_labels,
	no_hypothesis_template,
	no_candidate_labels,
	],
	outputs=output_textbox,
	)

	# Run the app
	demo.launch(width=.5)


	@torch.compile()
	def classification(
	doc: str,
	da_hypothesis_template: str,
	da_candidate_labels: str,
	sv_hypothesis_template: str,
	sv_candidate_labels: str,
	no_hypothesis_template: str,
	no_candidate_labels: str,
	) -> Dict[str, float]:
	"""Classify text into categories.

	Args:
	doc (str):
	Text to classify.
	da_hypothesis_template (str):
	Template for the hypothesis to be used for Danish classification.
	da_candidate_labels (str):
	Comma-separated list of candidate labels for Danish classification.
	sv_hypothesis_template (str):
	Template for the hypothesis to be used for Swedish classification.
	sv_candidate_labels (str):
	Comma-separated list of candidate labels for Swedish classification.
	no_hypothesis_template (str):
	Template for the hypothesis to be used for Norwegian classification.
	no_candidate_labels (str):
	Comma-separated list of candidate labels for Norwegian classification.

	Returns:
	dict of str to float:
	The predicted label and the confidence score.
	"""
	# Detect the language of the text
	language = detect_language(doc.replace('\n', ' ')).name

	# Set the hypothesis template and candidate labels based on the detected language
	if language == "sv":
	hypothesis_template = sv_hypothesis_template
	candidate_labels = re.split(r', *', sv_candidate_labels)
	elif language == "no":
	hypothesis_template = no_hypothesis_template
	candidate_labels = re.split(r', *', no_candidate_labels)
	else:
	hypothesis_template = da_hypothesis_template
	candidate_labels = re.split(r', *', da_candidate_labels)

	# Run the classifier on the text
	result = classifier(
	doc,
	candidate_labels=candidate_labels,
	hypothesis_template=hypothesis_template,
	)

	print(result)

	# Return the predicted label
	return {lbl: score for lbl, score in zip(result["labels"], result["scores"])}


	if __name__ == "__main__":
	main()