|
|
from transformers import pipeline, BartTokenizer, BartForSequenceClassification |
|
|
class ZeroShotClassifier: |
|
|
|
|
|
def __init__(self, model_name): |
|
|
self.model = self.create_model(model_name) |
|
|
self.model_name = model_name |
|
|
self.sentiment_labels = ["Positive", "Negative", "Neutral"] |
|
|
self.intention_labels = ["Inquire", "Inform", "Payment", "Price", "Trade In", "Discount", "Complaint", "Approve", "Selling", "Confusion", "Change Package", "Upgrade", "Purchase", "Help"] |
|
|
self.labels = self.sentiment_labels + self.intention_labels |
|
|
|
|
|
def create_model(self, model_name): |
|
|
|
|
|
tokenizer = BartTokenizer.from_pretrained(model_name) |
|
|
model = BartForSequenceClassification.from_pretrained(model_name) |
|
|
classifier = pipeline("zero-shot-classification", model=model, tokenizer=tokenizer) |
|
|
return classifier |
|
|
|
|
|
def analyze_text(self, text): |
|
|
results = list(self.model(text, self.labels)['labels']) |
|
|
i = 0 |
|
|
sentiment = None |
|
|
intention = None |
|
|
while (sentiment is None) or (intention is None): |
|
|
if results[i] in self.sentiment_labels: |
|
|
|
|
|
sentiment = results[i] |
|
|
if results[i] in self.intention_labels: |
|
|
|
|
|
intention = results[i] |
|
|
i += 1 |
|
|
return {"sentiment": sentiment, "intention": intention} |