diff --git a/tests/unit/test_youtube_extraction.py b/tests/unit/test_youtube_extraction.py new file mode 100644 index 0000000..3d2d90c --- /dev/null +++ b/tests/unit/test_youtube_extraction.py @@ -0,0 +1,34 @@ +"""Pure-helper tests for the claim extractor — the anti-hallucination quote match + chunking. +The Ollama call itself is integration-tested live (not mocked here).""" + +from youtube_claims.extraction import _chunks, _norm, _normalize, _quote_supported + + +def test_norm_collapses_whitespace(): + assert _norm(" A\nB c ") == "a b c" + + +def test_quote_supported_matches_across_segment_joins(): + # the caption text is fragmented (joined with spaces); a quote spanning the join still matches + ntext = _norm("as it broke through 12\ni took 3000 off the table at 12 and 13 cents") + assert _quote_supported("as it broke through 12 I took 3000 off the table", ntext) + + +def test_quote_supported_rejects_hallucination(): + ntext = _norm("the trader discussed a reverse split squeeze") + assert not _quote_supported("AAPL will hit 300 next week", ntext) + assert not _quote_supported("", ntext) + + +def test_normalize_drops_claim_without_verbatim_quote(): + assert _normalize({"asset_ticker": "aapl", "verbatim_quote": ""}, 0, 6) is None + ok = _normalize({"asset_ticker": "aapl", "verbatim_quote": "aapl broke out"}, 0, 6) + assert ok and ok["asset_ticker"] == "AAPL" and ok["timestamp_start"] == 0 + + +def test_chunks_group_segments_with_time_span(): + segs = [{"start": float(i), "end": float(i + 1), "text": "x" * 100} for i in range(200)] + wins = list(_chunks(segs, max_chars=1000)) + assert len(wins) > 1 + buf, start, end = wins[0] + assert start == 0.0 and end is not None and len(buf) >= 1 diff --git a/youtube_claims/extraction.py b/youtube_claims/extraction.py index e8523a6..d806bb2 100644 --- a/youtube_claims/extraction.py +++ b/youtube_claims/extraction.py @@ -5,10 +5,25 @@ verbatim_quote is dropped: that quote is the chain of custody the reasoning mode against ASR error.""" import json +import re import urllib.request from youtube_claims import config + +def _norm(s: str) -> str: + """Collapse whitespace + lowercase so a verbatim quote still matches across the segment + joins and the punctuation-poor text of auto-captions.""" + return re.sub(r"\s+", " ", s or "").strip().lower() + + +def _quote_supported(verbatim_quote: str, ntext: str) -> bool: + """Anti-hallucination: is the quote actually present in the (already whitespace-normalized) + source window? Matches a ~40-char prefix to tolerate minor ASR word-boundary noise while + still rejecting a fabricated quote.""" + nvq = _norm(verbatim_quote) + return bool(nvq) and (nvq[:40] in ntext or nvq in ntext) + _PROMPT = """You are a precise information extractor. From the TRANSCRIPT below, extract every distinct market/trading CLAIM about a specific asset (stock, crypto, ETF, index, commodity, fx). @@ -94,16 +109,16 @@ def extract_claims(segments: list[dict], watchlist: list[str] | None = None, wl = ", ".join(watchlist if watchlist is not None else config.watchlist()) or "(none specified)" out: list[dict] = [] for win, w_start, w_end in _chunks(segments): - text = "\n".join(t for _, t in win) + text = " ".join(t for _, t in win) + ntext = _norm(text) try: raw = _ollama_generate(_PROMPT.format(watchlist=wl, transcript=text), model, host) data = json.loads(raw) except Exception: # noqa: BLE001 — a bad chunk shouldn't kill the whole video continue for c in (data.get("claims") or []): - vq = (c.get("verbatim_quote") or "").strip() # keep only claims whose quote is actually present in THIS window (anti-hallucination) - if vq and vq[:60].lower() in text.lower(): + if _quote_supported(c.get("verbatim_quote") or "", ntext): nc = _normalize(c, w_start, w_end) if nc: out.append(nc) diff --git a/youtube_claims/momentum-transcripts.service b/youtube_claims/momentum-transcripts.service new file mode 100644 index 0000000..2c9bec4 --- /dev/null +++ b/youtube_claims/momentum-transcripts.service @@ -0,0 +1,19 @@ +[Unit] +Description=Momentum YouTube transcript -> claim-extraction pipeline +StartLimitIntervalSec=0 + +[Service] +Type=simple +WorkingDirectory=/mnt/c/code/momentum +# secrets (DATABASE_URL, YOUTUBE_API_KEY, YT_WATCHLIST) load from the repo's gitignored .env +# via youtube_claims.config (python-dotenv). Whisper is option A (CPU) by default; flip to the +# LAN GPU (option B) with YT_WHISPER_MODEL=large-v3 YT_WHISPER_DEVICE=cuda YT_WHISPER_COMPUTE_TYPE=float16. +ExecStart=/home/philip/.venvs/momentum/bin/python -u -m youtube_claims.run run +Restart=always +RestartSec=30 +StandardOutput=append:/mnt/c/code/momentum/data/transcripts_pipeline.log +StandardError=append:/mnt/c/code/momentum/data/transcripts_pipeline.log +SyslogIdentifier=momentum-transcripts + +[Install] +WantedBy=default.target