Add sentiment_engine with CryptoSentimentCalibrator fixes - improved keyword lists, lowered FinBERT threshold, added neutral handling
This commit is contained in:
Codex
@@ -0,0 +1,322 @@
|
||||
"""
|
||||
Comprehensive tests for SentimentEmotionAnalyzer with various edge cases.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import numpy as np
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from sentiment_engine.nlp.sentiment_emotion import (
|
||||
SentimentEmotionAnalyzer, ONNXSentimentModel, ONNXEmotionModel,
|
||||
CryptoSentimentCalibrator, MockTokenizer, MockSentimentModel
|
||||
)
|
||||
from sentiment_engine.schemas.processed import SentimentScores, EmotionScores
|
||||
|
||||
|
||||
class TestCryptoSentimentCalibrator:
|
||||
"""Tests for the crypto sentiment calibrator"""
|
||||
|
||||
def test_calibrate_no_flip_when_aligned(self):
|
||||
"""Should not flip when crypto and FinBERT signals align"""
|
||||
# Crypto bullish, FinBERT positive (bullish)
|
||||
probs = np.array([0.1, 0.2, 0.7]) # [neg, neu, pos]
|
||||
calibrated = CryptoSentimentCalibrator.calibrate("Bitcoin surges to new high", probs)
|
||||
np.testing.assert_array_almost_equal(calibrated, probs)
|
||||
|
||||
def test_calibrate_flip_bullish_crypto_bearish_finbert(self):
|
||||
"""Should flip when crypto says bullish but FinBERT says bearish"""
|
||||
probs = np.array([0.85, 0.1, 0.05]) # FinBERT: negative
|
||||
calibrated = CryptoSentimentCalibrator.calibrate("Bitcoin surges to $100k", probs)
|
||||
# Should flip: neg becomes pos
|
||||
assert calibrated[2] > calibrated[0] # pos > neg
|
||||
|
||||
def test_calibrate_flip_bearish_crypto_bullish_finbert(self):
|
||||
"""Should flip when crypto says bearish but FinBERT says bullish"""
|
||||
probs = np.array([0.05, 0.1, 0.85]) # FinBERT: positive
|
||||
calibrated = CryptoSentimentCalibrator.calibrate("Bitcoin crashes 50%", probs)
|
||||
# Should flip: pos becomes neg
|
||||
assert calibrated[0] > calibrated[2] # neg > pos
|
||||
|
||||
def test_calibrate_no_flip_neutral_crypto(self):
|
||||
"""Should not flip when crypto signal is neutral"""
|
||||
probs = np.array([0.3, 0.5, 0.2])
|
||||
calibrated = CryptoSentimentCalibrator.calibrate("BTC at $50k", probs)
|
||||
np.testing.assert_array_almost_equal(calibrated, probs)
|
||||
|
||||
def test_calibrate_preserves_probabilities_sum(self):
|
||||
"""Calibrated probabilities should sum to 1"""
|
||||
probs = np.array([0.85, 0.1, 0.05])
|
||||
calibrated = CryptoSentimentCalibrator.calibrate("Bitcoin surges", probs)
|
||||
assert abs(calibrated.sum() - 1.0) < 0.001
|
||||
|
||||
def test_get_crypto_signal_bullish(self):
|
||||
"""Should detect bullish signal from keywords"""
|
||||
signal = CryptoSentimentCalibrator._get_crypto_signal("Bitcoin surges to new ATH")
|
||||
assert signal == "bullish"
|
||||
|
||||
def test_get_crypto_signal_bearish(self):
|
||||
"""Should detect bearish signal from keywords"""
|
||||
signal = CryptoSentimentCalibrator._get_crypto_signal("Bitcoin crashes hard")
|
||||
assert signal == "bearish"
|
||||
|
||||
def test_get_crypto_signal_neutral(self):
|
||||
"""Should detect neutral when no strong signals"""
|
||||
signal = CryptoSentimentCalibrator._get_crypto_signal("BTC at $50k")
|
||||
assert signal == "neutral"
|
||||
|
||||
def test_get_finbert_signal_bullish(self):
|
||||
"""Should detect FinBERT bullish from probs"""
|
||||
probs = np.array([0.1, 0.2, 0.7])
|
||||
signal = CryptoSentimentCalibrator._get_finbert_signal(probs)
|
||||
assert signal == "bullish"
|
||||
|
||||
def test_get_finbert_signal_bearish(self):
|
||||
"""Should detect FinBERT bearish from probs"""
|
||||
probs = np.array([0.8, 0.15, 0.05])
|
||||
signal = CryptoSentimentCalibrator._get_finbert_signal(probs)
|
||||
assert signal == "bearish"
|
||||
|
||||
def test_get_finbert_signal_neutral(self):
|
||||
"""Should detect FinBERT neutral from probs"""
|
||||
probs = np.array([0.35, 0.4, 0.25])
|
||||
signal = CryptoSentimentCalibrator._get_finbert_signal(probs)
|
||||
assert signal == "neutral"
|
||||
|
||||
|
||||
class TestSentimentEmotionAnalyzer:
|
||||
"""Tests for SentimentEmotionAnalyzer"""
|
||||
|
||||
@pytest.fixture
|
||||
def analyzer(self):
|
||||
return SentimentEmotionAnalyzer()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_loads_model(self, analyzer):
|
||||
"""Should initialize and load model"""
|
||||
await analyzer.initialize()
|
||||
assert analyzer._model is not None
|
||||
assert analyzer._tokenizer is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_single_asset(self, analyzer):
|
||||
"""Should analyze sentiment for single asset"""
|
||||
await analyzer.initialize()
|
||||
|
||||
text = "Bitcoin surges to new all-time high!"
|
||||
asset_mentions = [{"asset_id": "BTC", "span": (0, 3)}]
|
||||
|
||||
sentiment_results, emotion_results = await analyzer.analyze(text, asset_mentions)
|
||||
|
||||
assert "BTC" in sentiment_results
|
||||
assert isinstance(sentiment_results["BTC"], SentimentScores)
|
||||
assert -1 <= sentiment_results["BTC"].polarity <= 1
|
||||
assert 0 <= sentiment_results["BTC"].confidence <= 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_multiple_assets(self, analyzer):
|
||||
"""Should analyze sentiment for multiple assets"""
|
||||
await analyzer.initialize()
|
||||
|
||||
text = "BTC and ETH both surge"
|
||||
asset_mentions = [
|
||||
{"asset_id": "BTC", "span": (0, 3)},
|
||||
{"asset_id": "ETH", "span": (8, 11)}
|
||||
]
|
||||
|
||||
sentiment_results, emotion_results = await analyzer.analyze(text, asset_mentions)
|
||||
|
||||
assert "BTC" in sentiment_results
|
||||
assert "ETH" in sentiment_results
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_empty_assets(self, analyzer):
|
||||
"""Should handle empty asset mentions"""
|
||||
await analyzer.initialize()
|
||||
|
||||
text = "Market is volatile"
|
||||
asset_mentions = []
|
||||
|
||||
sentiment_results, emotion_results = await analyzer.analyze(text, asset_mentions)
|
||||
|
||||
assert sentiment_results == {}
|
||||
assert emotion_results == {}
|
||||
|
||||
def test_heuristic_sentiment_bullish(self, analyzer):
|
||||
"""Heuristic should detect bullish sentiment"""
|
||||
text = "Bitcoin surges to new high! Bullish!"
|
||||
scores = analyzer._heuristic_sentiment(text)
|
||||
|
||||
assert scores.polarity > 0
|
||||
assert scores.positive_prob > scores.negative_prob
|
||||
|
||||
def test_heuristic_sentiment_bearish(self, analyzer):
|
||||
"""Heuristic should detect bearish sentiment"""
|
||||
text = "Bitcoin crashes hard! Panic selling!"
|
||||
scores = analyzer._heuristic_sentiment(text)
|
||||
|
||||
assert scores.polarity < 0
|
||||
assert scores.negative_prob > scores.positive_prob
|
||||
|
||||
def test_heuristic_sentiment_neutral(self, analyzer):
|
||||
"""Heuristic should detect neutral sentiment"""
|
||||
text = "BTC at $50,000, ETH at $3,000"
|
||||
scores = analyzer._heuristic_sentiment(text)
|
||||
|
||||
assert abs(scores.polarity) < 0.5
|
||||
|
||||
def test_heuristic_emotions_joy(self, analyzer):
|
||||
"""Heuristic should detect joy"""
|
||||
text = "Bitcoin mooning! Profit! Gains!"
|
||||
scores = analyzer._heuristic_emotions(text)
|
||||
|
||||
assert scores.joy > 0.5
|
||||
|
||||
def test_heuristic_emotions_fear(self, analyzer):
|
||||
"""Heuristic should detect fear"""
|
||||
text = "Crash! Panic! Liquidation! Fear!"
|
||||
scores = analyzer._heuristic_emotions(text)
|
||||
|
||||
assert scores.fear > 0.5
|
||||
|
||||
def test_heuristic_emotions_anger(self, analyzer):
|
||||
"""Heuristic should detect anger"""
|
||||
text = "Scam! Fraud! Rug pull! Unfair!"
|
||||
scores = analyzer._heuristic_emotions(text)
|
||||
|
||||
assert scores.anger > 0.5
|
||||
|
||||
def test_heuristic_emotions_greed(self, analyzer):
|
||||
"""Heuristic should detect greed"""
|
||||
text = "Buy buy buy! FOMO! YOLO! Leverage!"
|
||||
scores = analyzer._heuristic_emotions(text)
|
||||
|
||||
assert scores.greed > 0.5
|
||||
|
||||
def test_heuristic_emotions_sadness(self, analyzer):
|
||||
"""Heuristic should detect sadness"""
|
||||
text = "Lost everything. Rekt. Pain."
|
||||
scores = analyzer._heuristic_emotions(text)
|
||||
|
||||
assert scores.sadness > 0.5
|
||||
|
||||
def test_compute_intensity_high(self, analyzer):
|
||||
"""Should compute high intensity for emotional text"""
|
||||
text = "CRASH!!! BTC DUMPING!!!"
|
||||
intensity = analyzer.compute_intensity(text)
|
||||
|
||||
assert intensity > 0.5
|
||||
|
||||
def test_compute_intensity_low(self, analyzer):
|
||||
"""Should compute low intensity for neutral text"""
|
||||
text = "BTC at $50k"
|
||||
intensity = analyzer.compute_intensity(text)
|
||||
|
||||
assert intensity < 0.5
|
||||
|
||||
|
||||
class TestONNXSentimentModel:
|
||||
"""Tests for ONNXSentimentModel wrapper"""
|
||||
|
||||
def test_init_loads_session(self):
|
||||
"""Should load ONNX session"""
|
||||
with patch('onnxruntime.InferenceSession') as mock_session:
|
||||
mock_session.return_value.get_inputs.return_value = [
|
||||
MagicMock(name="input_ids"),
|
||||
MagicMock(name="attention_mask"),
|
||||
MagicMock(name="token_type_ids")
|
||||
]
|
||||
mock_session.return_value.get_outputs.return_value = [
|
||||
MagicMock(name="logits")
|
||||
]
|
||||
|
||||
with patch('transformers.AutoTokenizer.from_pretrained'):
|
||||
model = ONNXSentimentModel("path", "tokenizer_path")
|
||||
assert model.session is not None
|
||||
|
||||
def test_call_returns_logits(self):
|
||||
"""__call__ should return logits"""
|
||||
with patch('onnxruntime.InferenceSession') as mock_session:
|
||||
mock_session.return_value.get_inputs.return_value = [
|
||||
MagicMock(name="input_ids"),
|
||||
MagicMock(name="attention_mask"),
|
||||
MagicMock(name="token_type_ids")
|
||||
]
|
||||
mock_session.return_value.get_outputs.return_value = [MagicMock(name="logits")]
|
||||
mock_session.return_value.run.return_value = [np.array([[0.1, 0.2, 0.7]])]
|
||||
|
||||
with patch('transformers.AutoTokenizer.from_pretrained'):
|
||||
model = ONNXSentimentModel("path", "tokenizer_path")
|
||||
logits = model(
|
||||
np.ones((1, 10), dtype=np.int64),
|
||||
np.ones((1, 10), dtype=np.int64)
|
||||
)
|
||||
assert logits.shape == (1, 3)
|
||||
|
||||
|
||||
class TestONNXEmotionModel:
|
||||
"""Tests for ONNXEmotionModel wrapper"""
|
||||
|
||||
def test_init_loads_session(self):
|
||||
"""Should load ONNX session"""
|
||||
with patch('onnxruntime.InferenceSession') as mock_session:
|
||||
mock_session.return_value.get_inputs.return_value = [
|
||||
MagicMock(name="input_ids"),
|
||||
MagicMock(name="attention_mask")
|
||||
]
|
||||
mock_session.return_value.get_outputs.return_value = [MagicMock(name="logits")]
|
||||
|
||||
with patch('transformers.AutoTokenizer.from_pretrained'):
|
||||
model = ONNXEmotionModel("path", "tokenizer_path")
|
||||
assert model.session is not None
|
||||
|
||||
def test_call_returns_logits(self):
|
||||
"""__call__ should return logits"""
|
||||
with patch('onnxruntime.InferenceSession') as mock_session:
|
||||
mock_session.return_value.get_inputs.return_value = [
|
||||
MagicMock(name="input_ids"),
|
||||
MagicMock(name="attention_mask")
|
||||
]
|
||||
mock_session.return_value.get_outputs.return_value = [MagicMock(name="logits")]
|
||||
mock_session.return_value.run.return_value = [np.array([[0.1, 0.2, 0.3, 0.4, 0.0, 0.0]])]
|
||||
|
||||
with patch('transformers.AutoTokenizer.from_pretrained'):
|
||||
model = ONNXEmotionModel("path", "tokenizer_path")
|
||||
logits = model(
|
||||
np.ones((1, 10), dtype=np.int64),
|
||||
np.ones((1, 10), dtype=np.int64)
|
||||
)
|
||||
assert logits.shape == (1, 6)
|
||||
|
||||
|
||||
class TestMockComponents:
|
||||
"""Tests for mock components"""
|
||||
|
||||
def test_mock_tokenizer_returns_dict(self):
|
||||
"""MockTokenizer should return dict with required keys"""
|
||||
tokenizer = MockTokenizer()
|
||||
result = tokenizer("test text")
|
||||
|
||||
assert "input_ids" in result
|
||||
assert "attention_mask" in result
|
||||
assert "token_type_ids" in result
|
||||
|
||||
def test_mock_tokenizer_batch(self):
|
||||
"""MockTokenizer should handle batch input"""
|
||||
tokenizer = MockTokenizer()
|
||||
result = tokenizer(["text1", "text2"])
|
||||
|
||||
assert "input_ids" in result
|
||||
assert result["input_ids"].shape[0] == 2
|
||||
|
||||
def test_mock_sentiment_model(self):
|
||||
"""MockSentimentModel should return logits"""
|
||||
model = MockSentimentModel()
|
||||
result = model(input_ids=np.ones((2, 10)), attention_mask=np.ones((2, 10)))
|
||||
|
||||
assert hasattr(result, 'logits')
|
||||
assert result.logits.shape == (2, 3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Reference in New Issue
Block a user