Celes Renata
a72f336ad1
Multi-stage evidence-grounded inference architecture replacing the monolithic 9B model extraction pipeline. CPU-first specialist services handle routine extraction while the 9B vLLM model is preserved for semantic adjudication of ambiguous cases. Key components: - Capability-aware inference gateway (OpenAI-compatible + Ollama) - Endpoint registry with DB migrations and REST API - Sentence-aware document segmenter (property tests) - Deterministic financial parsing with offset integrity - Symbol resolution with ambiguity detection - Specialist service (GLiNER2, dynamic batching, K8s deployment) - Company-specific sentiment (FinBERT, calibration) - Retrieval-based novelty and duplicate detection - Confidence calibration pipeline - Deterministic routing engine (property tests) - 9B adjudication layer with VRAM gating - Stock-specific impact model (features, labels, baseline, trained) - Pipeline orchestrator (state machine, queues, leases, feature flags) - Bounded parallelism (async workers, semaphore, load shedding) - Observability (tracing, metrics, alerts) - Compatibility adapter (v3→v2 golden mapping tests) - Shadow/canary promotion framework - Active learning and fine-tuning pipeline Test results: 1,161 tests pass, ruff lint clean. All 282 spec tasks completed.
858 lines
30 KiB
Python
858 lines
30 KiB
Python
"""Tests for company-specific sentiment analysis.
|
|
|
|
Validates:
|
|
- Evidence grouping by company (including relations)
|
|
- FinBERT adapter returns valid probability distributions
|
|
- Mixed sentiment detection from evidence-group disagreement
|
|
- Non-mixed when evidence agrees
|
|
- Probability distributions sum to ~1.0
|
|
- Calibration passthrough
|
|
- SentimentScorer integration
|
|
- TextSentiment per-text scoring
|
|
- Aggregation module
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from services.intelligence_pipeline_v3.sentiment.aggregation import (
|
|
MIXED_DISAGREEMENT_THRESHOLD,
|
|
aggregate_evidence_sentiments,
|
|
)
|
|
from services.intelligence_pipeline_v3.sentiment.calibrator import SentimentCalibrator
|
|
from services.intelligence_pipeline_v3.sentiment.evidence_groups import build_evidence_groups
|
|
from services.intelligence_pipeline_v3.sentiment.finbert_adapter import FinBERTAdapter
|
|
from services.intelligence_pipeline_v3.sentiment.mixed_sentiment import (
|
|
DISAGREEMENT_THRESHOLD,
|
|
compute_mixed_sentiment,
|
|
)
|
|
from services.intelligence_pipeline_v3.sentiment.models import (
|
|
CompanySentimentResult,
|
|
EvidenceGroup,
|
|
SentimentBatchResult,
|
|
TextSentiment,
|
|
)
|
|
from services.intelligence_pipeline_v3.sentiment.sentiment_scorer import (
|
|
SentimentScorer,
|
|
)
|
|
|
|
|
|
class TestEvidenceGroups:
|
|
"""Test evidence grouping by company."""
|
|
|
|
def test_single_company_single_evidence(self):
|
|
entities = [{"company_id": "AAPL", "evidence_id": "ev1"}]
|
|
evidence_spans = {"ev1": "Apple reported strong earnings."}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
|
|
assert "AAPL" in groups
|
|
assert groups["AAPL"].company_id == "AAPL"
|
|
assert groups["AAPL"].evidence_ids == ["ev1"]
|
|
assert groups["AAPL"].texts == ["Apple reported strong earnings."]
|
|
|
|
def test_single_company_multiple_evidence(self):
|
|
entities = [
|
|
{"company_id": "AAPL", "evidence_id": "ev1"},
|
|
{"company_id": "AAPL", "evidence_id": "ev2"},
|
|
]
|
|
evidence_spans = {
|
|
"ev1": "Apple beat expectations.",
|
|
"ev2": "iPhone sales surged.",
|
|
}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
|
|
assert "AAPL" in groups
|
|
assert len(groups["AAPL"].evidence_ids) == 2
|
|
assert "ev1" in groups["AAPL"].evidence_ids
|
|
assert "ev2" in groups["AAPL"].evidence_ids
|
|
|
|
def test_multiple_companies(self):
|
|
entities = [
|
|
{"company_id": "AAPL", "evidence_id": "ev1"},
|
|
{"company_id": "GOOGL", "evidence_id": "ev2"},
|
|
]
|
|
evidence_spans = {
|
|
"ev1": "Apple gained market share.",
|
|
"ev2": "Google's ad revenue declined.",
|
|
}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
|
|
assert len(groups) == 2
|
|
assert "AAPL" in groups
|
|
assert "GOOGL" in groups
|
|
|
|
def test_shared_evidence_across_companies(self):
|
|
"""A span mentioning multiple companies should appear in both groups."""
|
|
entities = [
|
|
{"company_id": "AAPL", "evidence_id": "ev1"},
|
|
{"company_id": "GOOGL", "evidence_id": "ev1"},
|
|
]
|
|
evidence_spans = {"ev1": "Apple and Google both reported growth."}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
|
|
assert "AAPL" in groups
|
|
assert "GOOGL" in groups
|
|
assert "ev1" in groups["AAPL"].evidence_ids
|
|
assert "ev1" in groups["GOOGL"].evidence_ids
|
|
|
|
def test_entities_without_company_id_skipped(self):
|
|
entities = [
|
|
{"company_id": None, "evidence_id": "ev1"},
|
|
{"company_id": "AAPL", "evidence_id": "ev2"},
|
|
]
|
|
evidence_spans = {
|
|
"ev1": "Some generic text.",
|
|
"ev2": "Apple expanded.",
|
|
}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
|
|
assert len(groups) == 1
|
|
assert "AAPL" in groups
|
|
|
|
def test_missing_evidence_span_excluded(self):
|
|
"""Entity referencing non-existent evidence span is excluded."""
|
|
entities = [{"company_id": "AAPL", "evidence_id": "ev_missing"}]
|
|
evidence_spans = {"ev1": "Some text."}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
|
|
assert len(groups) == 0
|
|
|
|
def test_empty_inputs(self):
|
|
groups = build_evidence_groups([], {})
|
|
assert len(groups) == 0
|
|
|
|
def test_deduplicates_evidence_ids_per_company(self):
|
|
"""Same evidence_id referenced twice for same company shouldn't duplicate."""
|
|
entities = [
|
|
{"company_id": "AAPL", "evidence_id": "ev1"},
|
|
{"company_id": "AAPL", "evidence_id": "ev1"},
|
|
]
|
|
evidence_spans = {"ev1": "Apple news."}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
|
|
assert groups["AAPL"].evidence_ids == ["ev1"]
|
|
assert len(groups["AAPL"].texts) == 1
|
|
|
|
def test_relations_add_evidence_to_company(self):
|
|
"""Relations parameter links additional evidence to companies."""
|
|
entities = [{"company_id": "AAPL", "evidence_id": "ev1"}]
|
|
relations = [
|
|
{"company_id": "AAPL", "evidence_id": "ev2", "relation_type": "directly_affects"},
|
|
]
|
|
evidence_spans = {
|
|
"ev1": "Apple reported earnings.",
|
|
"ev2": "iPhone demand surged globally.",
|
|
}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans, relations=relations)
|
|
|
|
assert "AAPL" in groups
|
|
assert "ev1" in groups["AAPL"].evidence_ids
|
|
assert "ev2" in groups["AAPL"].evidence_ids
|
|
assert len(groups["AAPL"].evidence_ids) == 2
|
|
|
|
def test_relations_create_new_company_group(self):
|
|
"""Relations can create groups for companies not in entities."""
|
|
entities = [{"company_id": "AAPL", "evidence_id": "ev1"}]
|
|
relations = [
|
|
{"company_id": "GOOGL", "evidence_id": "ev2", "relation_type": "inferred_exposure"},
|
|
]
|
|
evidence_spans = {
|
|
"ev1": "Apple expanded.",
|
|
"ev2": "Google was affected.",
|
|
}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans, relations=relations)
|
|
|
|
assert "AAPL" in groups
|
|
assert "GOOGL" in groups
|
|
assert groups["GOOGL"].evidence_ids == ["ev2"]
|
|
|
|
def test_relations_none_skipped(self):
|
|
"""None relations parameter is handled gracefully."""
|
|
entities = [{"company_id": "AAPL", "evidence_id": "ev1"}]
|
|
evidence_spans = {"ev1": "Apple news."}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans, relations=None)
|
|
|
|
assert "AAPL" in groups
|
|
assert groups["AAPL"].evidence_ids == ["ev1"]
|
|
|
|
|
|
class TestFinBERTAdapter:
|
|
"""Test FinBERT adapter returns valid probability distributions."""
|
|
|
|
def setup_method(self):
|
|
self.adapter = FinBERTAdapter(test_mode=True)
|
|
|
|
def test_model_version_exposed(self):
|
|
assert self.adapter.model_version == "ProsusAI/finbert@v1.0"
|
|
assert self.adapter.model_name == "ProsusAI/finbert"
|
|
|
|
def test_empty_input(self):
|
|
result = self.adapter.classify([])
|
|
assert result == []
|
|
|
|
def test_positive_text(self):
|
|
result = self.adapter.classify(["Company reported strong profit growth."])
|
|
assert len(result) == 1
|
|
pos, neg, neu = result[0]
|
|
assert pos > neg
|
|
assert pos > neu
|
|
assert abs(pos + neg + neu - 1.0) < 1e-6
|
|
|
|
def test_negative_text(self):
|
|
result = self.adapter.classify(["Revenue declined sharply amid weak demand."])
|
|
assert len(result) == 1
|
|
pos, neg, neu = result[0]
|
|
assert neg > pos
|
|
assert neg > neu
|
|
|
|
def test_neutral_text(self):
|
|
result = self.adapter.classify(["The company held its annual general meeting today."])
|
|
assert len(result) == 1
|
|
pos, neg, neu = result[0]
|
|
assert neu > pos
|
|
assert neu > neg
|
|
|
|
def test_mixed_keywords_text(self):
|
|
result = self.adapter.classify(["Revenue growth was strong but the decline in margins hurt"])
|
|
assert len(result) == 1
|
|
pos, neg, neu = result[0]
|
|
assert pos >= 0.3
|
|
assert neg >= 0.3
|
|
|
|
def test_batch_classification(self):
|
|
texts = [
|
|
"Earnings beat expectations.",
|
|
"Stock plunged on weak results.",
|
|
"Board met to discuss routine matters.",
|
|
]
|
|
results = self.adapter.classify(texts)
|
|
assert len(results) == 3
|
|
assert results[0][0] > results[0][1]
|
|
assert results[1][1] > results[1][0]
|
|
assert results[2][2] > results[2][0]
|
|
assert results[2][2] > results[2][1]
|
|
|
|
def test_probabilities_sum_to_one(self):
|
|
texts = ["Strong growth.", "Major loss.", "Neutral report."]
|
|
results = self.adapter.classify(texts)
|
|
for pos, neg, neu in results:
|
|
assert abs(pos + neg + neu - 1.0) < 1e-6
|
|
assert pos >= 0.0
|
|
assert neg >= 0.0
|
|
assert neu >= 0.0
|
|
|
|
|
|
class TestAggregation:
|
|
"""Test aggregate_evidence_sentiments from the aggregation module."""
|
|
|
|
def test_single_positive_text(self):
|
|
scores = [TextSentiment(evidence_id="ev1", positive_prob=0.8, negative_prob=0.1, neutral_prob=0.1)]
|
|
result = aggregate_evidence_sentiments("AAPL", scores, "test_model")
|
|
|
|
assert result.label == "positive"
|
|
assert result.company_id == "AAPL"
|
|
assert result.is_mixed is False
|
|
assert len(result.per_text_scores) == 1
|
|
assert result.per_text_scores[0].evidence_id == "ev1"
|
|
|
|
def test_single_negative_text(self):
|
|
scores = [TextSentiment(evidence_id="ev1", positive_prob=0.1, negative_prob=0.8, neutral_prob=0.1)]
|
|
result = aggregate_evidence_sentiments("GOOGL", scores, "test_model")
|
|
|
|
assert result.label == "negative"
|
|
assert result.is_mixed is False
|
|
|
|
def test_mixed_from_disagreeing_texts(self):
|
|
"""Two texts: one positive, one negative -> mixed."""
|
|
scores = [
|
|
TextSentiment(evidence_id="ev1", positive_prob=0.75, negative_prob=0.10, neutral_prob=0.15),
|
|
TextSentiment(evidence_id="ev2", positive_prob=0.10, negative_prob=0.75, neutral_prob=0.15),
|
|
]
|
|
result = aggregate_evidence_sentiments("TSLA", scores, "test_model")
|
|
|
|
assert result.label == "mixed"
|
|
assert result.is_mixed is True
|
|
assert result.positive_prob > 0.3
|
|
assert result.negative_prob > 0.3
|
|
|
|
def test_not_mixed_when_agreement(self):
|
|
"""Two positive texts should not trigger mixed."""
|
|
scores = [
|
|
TextSentiment(evidence_id="ev1", positive_prob=0.70, negative_prob=0.15, neutral_prob=0.15),
|
|
TextSentiment(evidence_id="ev2", positive_prob=0.65, negative_prob=0.20, neutral_prob=0.15),
|
|
]
|
|
result = aggregate_evidence_sentiments("AAPL", scores, "test_model")
|
|
|
|
assert result.label == "positive"
|
|
assert result.is_mixed is False
|
|
|
|
def test_empty_scores_neutral(self):
|
|
result = aggregate_evidence_sentiments("X", [], "test_model")
|
|
assert result.label == "neutral"
|
|
assert result.neutral_prob == 1.0
|
|
assert result.is_mixed is False
|
|
|
|
def test_probabilities_sum_to_one(self):
|
|
scores = [
|
|
TextSentiment(evidence_id="ev1", positive_prob=0.60, negative_prob=0.25, neutral_prob=0.15),
|
|
TextSentiment(evidence_id="ev2", positive_prob=0.30, negative_prob=0.50, neutral_prob=0.20),
|
|
TextSentiment(evidence_id="ev3", positive_prob=0.10, negative_prob=0.10, neutral_prob=0.80),
|
|
]
|
|
result = aggregate_evidence_sentiments("X", scores, "test_model")
|
|
total = result.positive_prob + result.negative_prob + result.neutral_prob
|
|
assert abs(total - 1.0) < 1e-4
|
|
|
|
def test_evidence_ids_preserved(self):
|
|
scores = [
|
|
TextSentiment(evidence_id="ev1", positive_prob=0.5, negative_prob=0.3, neutral_prob=0.2),
|
|
TextSentiment(evidence_id="ev2", positive_prob=0.4, negative_prob=0.4, neutral_prob=0.2),
|
|
]
|
|
result = aggregate_evidence_sentiments("AAPL", scores, "model_v1")
|
|
|
|
assert result.evidence_ids == ["ev1", "ev2"]
|
|
assert result.model_version == "model_v1"
|
|
assert result.calibration_version == "uncalibrated"
|
|
|
|
def test_disagreement_threshold_boundary(self):
|
|
"""Both max pos and max neg must be >= threshold for mixed."""
|
|
threshold = MIXED_DISAGREEMENT_THRESHOLD
|
|
scores = [
|
|
TextSentiment(evidence_id="ev1", positive_prob=threshold, negative_prob=0.05, neutral_prob=1.0 - threshold - 0.05),
|
|
TextSentiment(evidence_id="ev2", positive_prob=0.05, negative_prob=threshold, neutral_prob=1.0 - threshold - 0.05),
|
|
]
|
|
result = aggregate_evidence_sentiments("X", scores, "test")
|
|
assert result.is_mixed is True
|
|
assert result.label == "mixed"
|
|
|
|
def test_below_disagreement_threshold_not_mixed(self):
|
|
"""Below threshold should not be mixed."""
|
|
threshold = MIXED_DISAGREEMENT_THRESHOLD
|
|
scores = [
|
|
TextSentiment(evidence_id="ev1", positive_prob=threshold - 0.01, negative_prob=0.05, neutral_prob=1.0 - (threshold - 0.01) - 0.05),
|
|
TextSentiment(evidence_id="ev2", positive_prob=0.05, negative_prob=threshold - 0.01, neutral_prob=1.0 - (threshold - 0.01) - 0.05),
|
|
]
|
|
result = aggregate_evidence_sentiments("X", scores, "test")
|
|
assert result.is_mixed is False
|
|
|
|
|
|
class TestMixedSentiment:
|
|
"""Test mixed sentiment detection from evidence-group disagreement (legacy API)."""
|
|
|
|
def test_single_positive_group(self):
|
|
result = compute_mixed_sentiment(
|
|
company_id="AAPL",
|
|
group_results=[(0.75, 0.10, 0.15)],
|
|
evidence_ids=["ev1"],
|
|
model_version="ProsusAI/finbert@v1.0",
|
|
)
|
|
|
|
assert result.label == "positive"
|
|
assert result.company_id == "AAPL"
|
|
assert result.positive_prob > result.negative_prob
|
|
assert result.is_mixed is False
|
|
|
|
def test_single_negative_group(self):
|
|
result = compute_mixed_sentiment(
|
|
company_id="GOOGL",
|
|
group_results=[(0.10, 0.75, 0.15)],
|
|
evidence_ids=["ev1"],
|
|
model_version="ProsusAI/finbert@v1.0",
|
|
)
|
|
|
|
assert result.label == "negative"
|
|
assert result.negative_prob > result.positive_prob
|
|
|
|
def test_single_neutral_group(self):
|
|
result = compute_mixed_sentiment(
|
|
company_id="MSFT",
|
|
group_results=[(0.15, 0.15, 0.70)],
|
|
evidence_ids=["ev1"],
|
|
model_version="ProsusAI/finbert@v1.0",
|
|
)
|
|
|
|
assert result.label == "neutral"
|
|
assert result.neutral_prob > result.positive_prob
|
|
assert result.neutral_prob > result.negative_prob
|
|
|
|
def test_mixed_from_disagreeing_groups(self):
|
|
"""Two groups: one positive, one negative -> mixed."""
|
|
group_results = [
|
|
(0.75, 0.10, 0.15),
|
|
(0.10, 0.75, 0.15),
|
|
]
|
|
result = compute_mixed_sentiment(
|
|
company_id="TSLA",
|
|
group_results=group_results,
|
|
evidence_ids=["ev1", "ev2"],
|
|
model_version="ProsusAI/finbert@v1.0",
|
|
)
|
|
|
|
assert result.label == "mixed"
|
|
assert result.is_mixed is True
|
|
assert result.positive_prob > 0.3
|
|
assert result.negative_prob > 0.3
|
|
|
|
def test_no_mixed_when_agreement(self):
|
|
"""Two positive groups should not trigger mixed."""
|
|
group_results = [
|
|
(0.70, 0.15, 0.15),
|
|
(0.65, 0.20, 0.15),
|
|
]
|
|
result = compute_mixed_sentiment(
|
|
company_id="AAPL",
|
|
group_results=group_results,
|
|
evidence_ids=["ev1", "ev2"],
|
|
model_version="ProsusAI/finbert@v1.0",
|
|
)
|
|
|
|
assert result.label == "positive"
|
|
assert result.is_mixed is False
|
|
|
|
def test_disagreement_threshold_boundary(self):
|
|
"""Both max pos and max neg must be >= threshold for mixed."""
|
|
group_results = [
|
|
(DISAGREEMENT_THRESHOLD, 0.05, 0.65),
|
|
(0.05, DISAGREEMENT_THRESHOLD, 0.65),
|
|
]
|
|
result = compute_mixed_sentiment(
|
|
company_id="X",
|
|
group_results=group_results,
|
|
evidence_ids=["ev1", "ev2"],
|
|
model_version="test",
|
|
)
|
|
assert result.label == "mixed"
|
|
assert result.is_mixed is True
|
|
|
|
def test_below_disagreement_threshold(self):
|
|
"""Below threshold should not be mixed."""
|
|
group_results = [
|
|
(DISAGREEMENT_THRESHOLD - 0.01, 0.05, 0.66),
|
|
(0.05, DISAGREEMENT_THRESHOLD - 0.01, 0.66),
|
|
]
|
|
result = compute_mixed_sentiment(
|
|
company_id="X",
|
|
group_results=group_results,
|
|
evidence_ids=["ev1", "ev2"],
|
|
model_version="test",
|
|
)
|
|
assert result.label == "neutral"
|
|
assert result.is_mixed is False
|
|
|
|
def test_empty_group_results(self):
|
|
result = compute_mixed_sentiment(
|
|
company_id="AAPL",
|
|
group_results=[],
|
|
evidence_ids=[],
|
|
model_version="test",
|
|
)
|
|
assert result.label == "neutral"
|
|
assert result.neutral_prob == 1.0
|
|
assert result.is_mixed is False
|
|
|
|
def test_probabilities_sum_to_one(self):
|
|
group_results = [
|
|
(0.60, 0.25, 0.15),
|
|
(0.30, 0.50, 0.20),
|
|
(0.10, 0.10, 0.80),
|
|
]
|
|
result = compute_mixed_sentiment(
|
|
company_id="X",
|
|
group_results=group_results,
|
|
evidence_ids=["a", "b", "c"],
|
|
model_version="test",
|
|
)
|
|
total = result.positive_prob + result.negative_prob + result.neutral_prob
|
|
assert abs(total - 1.0) < 1e-4
|
|
|
|
def test_per_text_scores_preserved(self):
|
|
"""Legacy API now populates per_text_scores for provenance."""
|
|
group_results = [(0.75, 0.10, 0.15), (0.20, 0.60, 0.20)]
|
|
result = compute_mixed_sentiment(
|
|
company_id="X",
|
|
group_results=group_results,
|
|
evidence_ids=["ev1", "ev2"],
|
|
model_version="test",
|
|
)
|
|
assert len(result.per_text_scores) == 2
|
|
assert result.per_text_scores[0].evidence_id == "ev1"
|
|
assert result.per_text_scores[1].evidence_id == "ev2"
|
|
|
|
|
|
class TestSentimentScorer:
|
|
"""Test SentimentScorer integration (end-to-end scoring)."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_score_positive_evidence(self):
|
|
scorer = SentimentScorer()
|
|
group = EvidenceGroup(
|
|
company_id="AAPL",
|
|
evidence_ids=["ev1"],
|
|
texts=["Apple reported strong profit growth."],
|
|
)
|
|
result = await scorer.score(group)
|
|
|
|
assert result.company_id == "AAPL"
|
|
assert result.label == "positive"
|
|
assert result.positive_prob > result.negative_prob
|
|
assert len(result.per_text_scores) == 1
|
|
assert result.per_text_scores[0].evidence_id == "ev1"
|
|
assert result.model_version == "ProsusAI/finbert@v1.0"
|
|
assert result.calibration_version == "uncalibrated"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_score_negative_evidence(self):
|
|
scorer = SentimentScorer()
|
|
group = EvidenceGroup(
|
|
company_id="GOOGL",
|
|
evidence_ids=["ev1"],
|
|
texts=["Google experienced a sharp decline in revenue."],
|
|
)
|
|
result = await scorer.score(group)
|
|
|
|
assert result.label == "negative"
|
|
assert result.negative_prob > result.positive_prob
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_score_mixed_evidence(self):
|
|
"""Multiple texts with opposing sentiment triggers mixed."""
|
|
scorer = SentimentScorer()
|
|
group = EvidenceGroup(
|
|
company_id="TSLA",
|
|
evidence_ids=["ev1", "ev2"],
|
|
texts=[
|
|
"Tesla revenue growth exceeded expectations.",
|
|
"Tesla faces major decline in margins and weak demand.",
|
|
],
|
|
)
|
|
result = await scorer.score(group)
|
|
|
|
assert result.label == "mixed"
|
|
assert result.is_mixed is True
|
|
assert len(result.per_text_scores) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_score_batch(self):
|
|
scorer = SentimentScorer()
|
|
groups = {
|
|
"AAPL": EvidenceGroup(
|
|
company_id="AAPL",
|
|
evidence_ids=["ev1"],
|
|
texts=["Apple beat earnings estimates."],
|
|
),
|
|
"GOOGL": EvidenceGroup(
|
|
company_id="GOOGL",
|
|
evidence_ids=["ev2"],
|
|
texts=["Google saw weak ad revenue and decline in users"],
|
|
),
|
|
}
|
|
batch_result = await scorer.score_batch(groups)
|
|
|
|
assert len(batch_result.results) == 2
|
|
assert batch_result.model_version == "ProsusAI/finbert@v1.0"
|
|
assert batch_result.processing_time_ms >= 0
|
|
|
|
labels = {r.company_id: r.label for r in batch_result.results}
|
|
assert labels["AAPL"] == "positive"
|
|
assert labels["GOOGL"] == "negative"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_score_probability_distributions_sum_to_one(self):
|
|
scorer = SentimentScorer()
|
|
group = EvidenceGroup(
|
|
company_id="X",
|
|
evidence_ids=["ev1", "ev2", "ev3"],
|
|
texts=["Profit rose.", "Demand weakened.", "Board meeting held."],
|
|
)
|
|
result = await scorer.score(group)
|
|
|
|
# Overall probabilities sum to 1
|
|
total = result.positive_prob + result.negative_prob + result.neutral_prob
|
|
assert abs(total - 1.0) < 1e-4
|
|
|
|
# Per-text probabilities also sum to 1
|
|
for ts in result.per_text_scores:
|
|
text_total = ts.positive_prob + ts.negative_prob + ts.neutral_prob
|
|
assert abs(text_total - 1.0) < 1e-6
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_model_protocol(self):
|
|
"""SentimentScorer works with any model implementing SentimentModel."""
|
|
|
|
class MockModel:
|
|
@property
|
|
def model_version(self) -> str:
|
|
return "mock@v1"
|
|
|
|
def classify(self, texts: list[str]) -> list[tuple[float, float, float]]:
|
|
return [(0.5, 0.3, 0.2)] * len(texts)
|
|
|
|
scorer = SentimentScorer(model=MockModel())
|
|
group = EvidenceGroup(
|
|
company_id="X",
|
|
evidence_ids=["ev1"],
|
|
texts=["Any text."],
|
|
)
|
|
result = await scorer.score(group)
|
|
|
|
assert result.model_version == "mock@v1"
|
|
assert result.positive_prob > 0.4
|
|
|
|
|
|
class TestMultiCompanyOpposingSentiments:
|
|
"""Test that opposing sentiments for different companies produce separate records."""
|
|
|
|
def test_separate_records_for_opposing_companies(self):
|
|
"""Article with positive Apple news and negative Google news."""
|
|
entities = [
|
|
{"company_id": "AAPL", "evidence_id": "ev1"},
|
|
{"company_id": "GOOGL", "evidence_id": "ev2"},
|
|
]
|
|
evidence_spans = {
|
|
"ev1": "Apple reported record profit growth.",
|
|
"ev2": "Google faces a major decline in ad revenue.",
|
|
}
|
|
|
|
groups = build_evidence_groups(entities, evidence_spans)
|
|
adapter = FinBERTAdapter(test_mode=True)
|
|
|
|
results: list[CompanySentimentResult] = []
|
|
for company_id, group in groups.items():
|
|
probs = adapter.classify(group.texts)
|
|
result = compute_mixed_sentiment(
|
|
company_id=company_id,
|
|
group_results=probs,
|
|
evidence_ids=group.evidence_ids,
|
|
model_version=adapter.model_version,
|
|
)
|
|
results.append(result)
|
|
|
|
assert len(results) == 2
|
|
company_labels = {r.company_id: r.label for r in results}
|
|
|
|
assert company_labels["AAPL"] == "positive"
|
|
assert company_labels["GOOGL"] == "negative"
|
|
|
|
|
|
class TestCalibrator:
|
|
"""Test sentiment probability calibration passthrough and fitting."""
|
|
|
|
def test_uncalibrated_passthrough(self):
|
|
"""Unfitted calibrator should pass through raw probabilities."""
|
|
cal = SentimentCalibrator(method="isotonic")
|
|
assert not cal.is_fitted
|
|
assert cal.calibration_version == "uncalibrated"
|
|
|
|
raw = [0.6, 0.3, 0.1]
|
|
result = cal.calibrate(raw)
|
|
assert result == raw
|
|
|
|
def test_isotonic_fit_and_calibrate(self):
|
|
"""Fitted isotonic calibrator should transform probabilities."""
|
|
cal = SentimentCalibrator(method="isotonic")
|
|
|
|
raw_probs = [
|
|
[0.8, 0.1, 0.1],
|
|
[0.7, 0.2, 0.1],
|
|
[0.1, 0.8, 0.1],
|
|
[0.2, 0.7, 0.1],
|
|
[0.1, 0.1, 0.8],
|
|
[0.1, 0.2, 0.7],
|
|
[0.9, 0.05, 0.05],
|
|
[0.05, 0.9, 0.05],
|
|
[0.05, 0.05, 0.9],
|
|
[0.6, 0.3, 0.1],
|
|
]
|
|
true_labels = [0, 0, 1, 1, 2, 2, 0, 1, 2, 0]
|
|
|
|
cal.fit(raw_probs, true_labels, version="gold_v1")
|
|
|
|
assert cal.is_fitted
|
|
assert cal.calibration_version == "gold_v1"
|
|
|
|
result = cal.calibrate([0.7, 0.2, 0.1])
|
|
assert len(result) == 3
|
|
assert all(0.0 <= p <= 1.0 for p in result)
|
|
assert abs(sum(result) - 1.0) < 1e-6
|
|
|
|
def test_calibration_preserves_ordering(self):
|
|
"""Higher raw probabilities should map to higher calibrated values."""
|
|
cal = SentimentCalibrator(method="isotonic")
|
|
|
|
raw_probs = [
|
|
[0.9, 0.05, 0.05],
|
|
[0.8, 0.1, 0.1],
|
|
[0.7, 0.15, 0.15],
|
|
[0.6, 0.2, 0.2],
|
|
[0.3, 0.6, 0.1],
|
|
[0.2, 0.7, 0.1],
|
|
[0.1, 0.8, 0.1],
|
|
[0.1, 0.1, 0.8],
|
|
[0.15, 0.15, 0.7],
|
|
[0.2, 0.2, 0.6],
|
|
]
|
|
true_labels = [0, 0, 0, 0, 1, 1, 1, 2, 2, 2]
|
|
|
|
cal.fit(raw_probs, true_labels, version="test_v1")
|
|
|
|
low_pos = cal.calibrate([0.3, 0.5, 0.2])
|
|
high_pos = cal.calibrate([0.8, 0.1, 0.1])
|
|
|
|
assert high_pos[0] >= low_pos[0]
|
|
|
|
def test_platt_calibration(self):
|
|
"""Platt scaling should also produce valid probabilities."""
|
|
cal = SentimentCalibrator(method="platt")
|
|
|
|
raw_probs = [
|
|
[0.8, 0.1, 0.1],
|
|
[0.7, 0.2, 0.1],
|
|
[0.1, 0.8, 0.1],
|
|
[0.2, 0.7, 0.1],
|
|
[0.1, 0.1, 0.8],
|
|
[0.1, 0.2, 0.7],
|
|
[0.9, 0.05, 0.05],
|
|
[0.05, 0.9, 0.05],
|
|
[0.05, 0.05, 0.9],
|
|
[0.6, 0.3, 0.1],
|
|
]
|
|
true_labels = [0, 0, 1, 1, 2, 2, 0, 1, 2, 0]
|
|
|
|
cal.fit(raw_probs, true_labels, version="platt_v1")
|
|
assert cal.is_fitted
|
|
|
|
result = cal.calibrate([0.6, 0.3, 0.1])
|
|
assert len(result) == 3
|
|
assert all(0.0 <= p <= 1.0 for p in result)
|
|
assert abs(sum(result) - 1.0) < 1e-6
|
|
|
|
def test_batch_calibrate(self):
|
|
"""Batch calibration should produce consistent results."""
|
|
cal = SentimentCalibrator(method="isotonic")
|
|
|
|
raw_probs = [
|
|
[0.9, 0.05, 0.05],
|
|
[0.1, 0.8, 0.1],
|
|
[0.1, 0.1, 0.8],
|
|
[0.7, 0.2, 0.1],
|
|
[0.2, 0.7, 0.1],
|
|
[0.2, 0.1, 0.7],
|
|
]
|
|
true_labels = [0, 1, 2, 0, 1, 2]
|
|
|
|
cal.fit(raw_probs, true_labels, version="batch_v1")
|
|
|
|
batch = [[0.7, 0.2, 0.1], [0.2, 0.7, 0.1]]
|
|
results = cal.calibrate_batch(batch)
|
|
|
|
assert len(results) == 2
|
|
for r in results:
|
|
assert abs(sum(r) - 1.0) < 1e-6
|
|
|
|
def test_fit_validation_errors(self):
|
|
cal = SentimentCalibrator()
|
|
|
|
with pytest.raises(ValueError):
|
|
cal.fit([], [])
|
|
|
|
with pytest.raises(ValueError):
|
|
cal.fit([[0.5, 0.3, 0.2]], [0, 1]) # Length mismatch
|
|
|
|
|
|
class TestModels:
|
|
"""Test data model validation."""
|
|
|
|
def test_evidence_group_requires_non_empty_ids(self):
|
|
with pytest.raises(ValueError):
|
|
EvidenceGroup(company_id="AAPL", evidence_ids=[], texts=["test"])
|
|
|
|
def test_evidence_group_requires_non_empty_texts(self):
|
|
with pytest.raises(ValueError):
|
|
EvidenceGroup(company_id="AAPL", evidence_ids=["ev1"], texts=[])
|
|
|
|
def test_company_sentiment_result_valid_labels(self):
|
|
for label in ("positive", "negative", "neutral", "mixed"):
|
|
result = CompanySentimentResult(
|
|
company_id="X",
|
|
label=label,
|
|
positive_prob=0.33,
|
|
negative_prob=0.33,
|
|
neutral_prob=0.34,
|
|
evidence_ids=["ev1"],
|
|
model_version="test",
|
|
)
|
|
assert result.label == label
|
|
|
|
def test_company_sentiment_result_invalid_label(self):
|
|
with pytest.raises(ValueError):
|
|
CompanySentimentResult(
|
|
company_id="X",
|
|
label="very_positive",
|
|
positive_prob=0.8,
|
|
negative_prob=0.1,
|
|
neutral_prob=0.1,
|
|
evidence_ids=["ev1"],
|
|
model_version="test",
|
|
)
|
|
|
|
def test_sentiment_batch_result(self):
|
|
result = SentimentBatchResult(
|
|
results=[
|
|
CompanySentimentResult(
|
|
company_id="AAPL",
|
|
label="positive",
|
|
positive_prob=0.8,
|
|
negative_prob=0.1,
|
|
neutral_prob=0.1,
|
|
evidence_ids=["ev1"],
|
|
model_version="test",
|
|
)
|
|
],
|
|
model_version="test",
|
|
processing_time_ms=150,
|
|
)
|
|
assert len(result.results) == 1
|
|
assert result.processing_time_ms == 150
|
|
|
|
def test_text_sentiment_model(self):
|
|
ts = TextSentiment(
|
|
evidence_id="ev1",
|
|
positive_prob=0.7,
|
|
negative_prob=0.2,
|
|
neutral_prob=0.1,
|
|
)
|
|
assert ts.evidence_id == "ev1"
|
|
assert ts.dominant_label == "positive"
|
|
assert abs(ts.positive_prob + ts.negative_prob + ts.neutral_prob - 1.0) < 1e-6
|
|
|
|
def test_text_sentiment_dominant_negative(self):
|
|
ts = TextSentiment(evidence_id="ev1", positive_prob=0.1, negative_prob=0.7, neutral_prob=0.2)
|
|
assert ts.dominant_label == "negative"
|
|
|
|
def test_text_sentiment_dominant_neutral(self):
|
|
ts = TextSentiment(evidence_id="ev1", positive_prob=0.1, negative_prob=0.2, neutral_prob=0.7)
|
|
assert ts.dominant_label == "neutral"
|
|
|
|
def test_company_sentiment_result_is_mixed_field(self):
|
|
result = CompanySentimentResult(
|
|
company_id="X",
|
|
label="mixed",
|
|
positive_prob=0.4,
|
|
negative_prob=0.4,
|
|
neutral_prob=0.2,
|
|
evidence_ids=["ev1", "ev2"],
|
|
is_mixed=True,
|
|
model_version="test",
|
|
)
|
|
assert result.is_mixed is True
|