datasetsANDmodels's picture
Update usage.py
7f64a87 verified
import sys
from transformers import pipeline
# Define candidate labels for classification
candidate_labels_spam = ['Spam', 'not Spam']
candidate_labels_urgent = ['Urgent', 'not Urgent']
model="SpamUrgencyDetection"
clf = pipeline("zero-shot-classification", model=model) 32
def predict(text):
p_spam = clf(text, candidate_labels_spam)["labels"][0]
p_urgent = clf(text, candidate_labels_urgent)["labels"][0]
return p_spam,p_urgent
import pandas as pd
df = pd.read_csv("test.csv")
texts=df["text"]
for i in range( len(texts)):
print(texts[i],predict(texts[i]))