Crypto-Analyst / sentiment_tools.py
menikev's picture
Update sentiment_tools.py
9b28320 verified
raw
history blame contribute delete
853 Bytes
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"