Update sentiment_tools.py

sentiment_tools.py

@@ -1,15 +1,22 @@
-from langchain.tools import tool
-from textblob import TextBlob
-
-class SentimentTools:
-    @tool("Analyze Sentiment")
-    def analyze_sentiment(text):
-        """Analyzes the sentiment of a given text."""
-        blob = TextBlob(text)
-        sentiment = blob.sentiment.polarity
-        if sentiment > 0:
-            return "Positive"
-        elif sentiment < 0:
-            return "Negative"
-        else:
+from langchain.tools import tool
+from transformers import pipeline
+
+class SentimentTools:
+    def __init__(self):
+        self.sentiment_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
+
+    @tool("Analyze Sentiment")
+    def analyze_sentiment(self, text):
+        """Analyzes the sentiment of a given text."""
+        # Limit the input text to 512 tokens to ensure faster processing
+        truncated_text = ' '.join(text.split()[:512])
+        result = self.sentiment_pipeline(truncated_text)[0]
+        sentiment = result['label']
+        score = result['score']
+
+        if sentiment == 'POSITIVE':
+            return "Positive" if score > 0.6 else "Neutral"
+        elif sentiment == 'NEGATIVE':
+            return "Negative" if score > 0.6 else "Neutral"
+        else:
             return "Neutral"