Text Classification
demo1 / sample.py
jeonseonjin's picture
commit
f9da252 verified
from transformers import pipeline
def classify_text(email):
"""
Use Facebook BART model to classify an email into "spam" or "not spam"
Args:
email (str): The email to classify
Returns:
str: The classification of the email ("spam" or "not spam")
"""
# Load the BART model for text classification
classifier = pipeline(task="zero-shot-classification", model="facebook/bart-large-mnli")
# Classify the email
labels = ['spam','not spam']
template = 'This email is {}.'
result = classifier(email, labels)
# Get the label with the highest score
label = result['labels'][0]
return label
classify_text('hi I am marketer, we have good product for your good life')