jatinmehra commited on
Commit
28ff814
·
1 Parent(s): 6083539

implement NegaBot model for tweet sentiment classification with logging and prediction capabilities

Browse files
Files changed (1) hide show
  1. model.py +126 -0
model.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ NegaBot Model - Tweet Sentiment Classification
3
+ Uses the SmolLM 360M V2 model for product criticism detection
4
+ """
5
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
6
+ import torch
7
+ import logging
8
+
9
+ # Configure logging
10
+ logging.basicConfig(level=logging.INFO)
11
+ logger = logging.getLogger(__name__)
12
+
13
+ class NegaBotModel:
14
+ def __init__(self, model_name="jatinmehra/NegaBot-Product-Criticism-Catcher"):
15
+ """
16
+ Initialize the NegaBot model for sentiment classification
17
+
18
+ Args:
19
+ model_name (str): HuggingFace model identifier
20
+ """
21
+ self.model_name = model_name
22
+ self.model = None
23
+ self.tokenizer = None
24
+ self.load_model()
25
+
26
+ def load_model(self):
27
+ """Load the model and tokenizer from HuggingFace"""
28
+ try:
29
+ logger.info(f"Loading model: {self.model_name}")
30
+ self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
31
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
32
+
33
+ # Set model to evaluation mode
34
+ self.model.eval()
35
+ logger.info("Model loaded successfully")
36
+
37
+ except Exception as e:
38
+ logger.error(f"Error loading model: {str(e)}")
39
+ raise e
40
+
41
+ def predict(self, text: str) -> dict:
42
+ """
43
+ Predict sentiment for a given text
44
+
45
+ Args:
46
+ text (str): Input text to classify
47
+
48
+ Returns:
49
+ dict: Prediction result with sentiment and confidence
50
+ """
51
+ try:
52
+ # Tokenize input text
53
+ inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
54
+
55
+ # Get model predictions
56
+ with torch.no_grad():
57
+ outputs = self.model(**inputs)
58
+ logits = outputs.logits
59
+
60
+ # Apply softmax to get probabilities
61
+ probabilities = torch.softmax(logits, dim=1)
62
+ predicted_class = torch.argmax(logits, dim=1).item()
63
+ confidence = probabilities[0][predicted_class].item()
64
+
65
+ # Map prediction to sentiment
66
+ sentiment = "Negative" if predicted_class == 1 else "Positive"
67
+
68
+ return {
69
+ "text": text,
70
+ "sentiment": sentiment,
71
+ "confidence": round(confidence, 4),
72
+ "predicted_class": predicted_class,
73
+ "probabilities": {
74
+ "positive": round(probabilities[0][0].item(), 4),
75
+ "negative": round(probabilities[0][1].item(), 4)
76
+ }
77
+ }
78
+
79
+ except Exception as e:
80
+ logger.error(f"Error during prediction: {str(e)}")
81
+ raise e
82
+
83
+ def batch_predict(self, texts: list) -> list:
84
+ """
85
+ Predict sentiment for multiple texts
86
+
87
+ Args:
88
+ texts (list): List of texts to classify
89
+
90
+ Returns:
91
+ list: List of prediction results
92
+ """
93
+ results = []
94
+ for text in texts:
95
+ results.append(self.predict(text))
96
+ return results
97
+
98
+ # Global model instance (singleton pattern)
99
+ _model_instance = None
100
+
101
+ def get_model():
102
+ """Get the global model instance"""
103
+ global _model_instance
104
+ if _model_instance is None:
105
+ _model_instance = NegaBotModel()
106
+ return _model_instance
107
+
108
+ if __name__ == "__main__":
109
+ # Test the model
110
+ model = NegaBotModel()
111
+
112
+ test_texts = [
113
+ "This product is awful and broke within a week!",
114
+ "Amazing quality, highly recommend this product!",
115
+ "The service was okay, nothing special.",
116
+ "Terrible customer support, waste of money!"
117
+ ]
118
+
119
+ print("Testing NegaBot Model:")
120
+ print("=" * 50)
121
+
122
+ for text in test_texts:
123
+ result = model.predict(text)
124
+ print(f"Text: {text}")
125
+ print(f"Sentiment: {result['sentiment']} (Confidence: {result['confidence']:.2%})")
126
+ print("-" * 30)