File size: 1,075 Bytes
afbf5e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | from transformers import pipeline, BartTokenizer, BartForSequenceClassification
class ZeroShotClassifier:
def __init__(self, model_name, sentiment_labels, intention_labels):
self.model = self.create_model(model_name)
self.model_name = model_name
self.sentiment_labels = sentiment_labels
self.intention_labels = intention_labels
def create_model(self, model_name):
# Create Model
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):
# Sentiment analysis
sentiment_result = self.model(text, self.sentiment_labels)
sentiment = sentiment_result["labels"][0]
# Intention analysis
intention_result = self.model(text, self.intention_labels)
intention = intention_result["labels"][0]
return {"sentiment": sentiment, "intention": intention} |