mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-16 23:01:06 +02:00
refactor: streamline TikTok and Instagram scraping logic by removing search_queries and enhancing documentation for clarity
This commit is contained in:
parent
e8b3692b54
commit
2b018c4474
111 changed files with 1800 additions and 1580 deletions
|
|
@ -14,11 +14,11 @@ Answer the delegated question from live TikTok data gathered with your verb, com
|
|||
</available_tools>
|
||||
|
||||
<playbook>
|
||||
- Finding videos on a topic: prefer `tiktok_scrape` with `hashtags` (no leading '#') or a direct TikTok URL in `urls` (fastest). `search_queries` also finds videos on a topic, but it is Google-backed and slow, so start with **at most 3** distinct queries and only add more if the first round returns nothing significant — never batch many phrasing variants of the same intent.
|
||||
- Finding videos on a topic: call `tiktok_scrape` with `hashtags` (no leading '#'), or pass a TikTok URL in `urls`. There is no keyword-video search — use hashtags or a video URL.
|
||||
- Scraping a specific video, profile, hashtag, or search page: pass its TikTok URL in `urls`.
|
||||
- Profiles: a creator's `profiles` feed returns the account's metadata (name, followers, bio, verification) reliably, but its video list is often withheld by TikTok — treat an empty video list as a known limit, not a failure to retry endlessly. Prefer `hashtags` or a direct video URL for videos.
|
||||
- Comments on a video: call `tiktok_comments` with the video URL(s) in `video_urls`.
|
||||
- Finding accounts by keyword: call `tiktok_user_search` with `queries` — that is the path for accounts. Use `search_queries` on `tiktok_scrape` only when you want videos, not accounts.
|
||||
- Finding accounts by keyword: call `tiktok_user_search` with `queries` — that is the path for accounts.
|
||||
- "What's trending now": call `tiktok_trending` (no query needed); set `max_items` for how many.
|
||||
- Controlling volume: use `max_items` for the total cap and `results_per_page` per target (per-verb equivalents: `comments_per_video`, `results_per_query`).
|
||||
- Requested counts: `max_items` defaults low — when the task asks for N items, set `max_items` (and the per-target count) above N. A call that caps below the target can never satisfy it.
|
||||
|
|
|
|||
|
|
@ -54,9 +54,7 @@ class DetailsInput(BaseModel):
|
|||
@model_validator(mode="after")
|
||||
def _exactly_one_source(self) -> DetailsInput:
|
||||
if not self.urls and not self.search_queries:
|
||||
raise ValueError(
|
||||
"Provide at least one of 'urls' or 'search_queries'."
|
||||
)
|
||||
raise ValueError("Provide at least one of 'urls' or 'search_queries'.")
|
||||
if self.urls and self.search_queries:
|
||||
raise ValueError(
|
||||
"Provide 'urls' OR 'search_queries', not both (they cannot be combined)."
|
||||
|
|
|
|||
|
|
@ -77,9 +77,7 @@ class ScrapeInput(BaseModel):
|
|||
@model_validator(mode="after")
|
||||
def _exactly_one_source(self) -> ScrapeInput:
|
||||
if not self.urls and not self.search_queries:
|
||||
raise ValueError(
|
||||
"Provide at least one of 'urls' or 'search_queries'."
|
||||
)
|
||||
raise ValueError("Provide at least one of 'urls' or 'search_queries'.")
|
||||
if self.urls and self.search_queries:
|
||||
raise ValueError(
|
||||
"Provide 'urls' OR 'search_queries', not both (they cannot be combined)."
|
||||
|
|
|
|||
|
|
@ -10,9 +10,8 @@ from app.capabilities.tiktok.scrape.schemas import ScrapeInput, ScrapeOutput
|
|||
TIKTOK_SCRAPE = Capability(
|
||||
name="tiktok.scrape",
|
||||
description=(
|
||||
"Scrape public TikTok videos. Use urls, profiles, hashtags, or "
|
||||
"search_queries (search_queries are resolved via Google to public "
|
||||
"videos; for accounts by keyword use tiktok.user_search)."
|
||||
"Scrape public TikTok videos. Use urls, profiles, or hashtags. To find "
|
||||
"accounts by keyword, use tiktok.user_search."
|
||||
),
|
||||
input_schema=ScrapeInput,
|
||||
output_schema=ScrapeOutput,
|
||||
|
|
|
|||
|
|
@ -26,7 +26,6 @@ def build_scrape_executor(scrape_fn: ScrapeFn | None = None) -> Executor:
|
|||
startUrls=[{"url": url} for url in payload.urls],
|
||||
profiles=payload.profiles,
|
||||
hashtags=payload.hashtags,
|
||||
searchQueries=payload.search_queries,
|
||||
resultsPerPage=payload.results_per_page,
|
||||
)
|
||||
emit_progress(
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ A lean, agent-friendly surface over ``TikTokScrapeInput``
|
|||
(``app/proprietary/platforms/tiktok``). The executor maps this to the full
|
||||
scraper input; the scraper's ``TikTokVideoItem`` is reused verbatim as the
|
||||
output element. Any TikTok URL kind (video, profile, hashtag, search) goes in
|
||||
``urls``; ``profiles``/``hashtags``/``search_queries`` are typed shortcuts.
|
||||
``urls``; ``profiles``/``hashtags`` are typed shortcuts. Keyword search is not a
|
||||
video source here — use ``tiktok.user_search`` to find accounts by keyword.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -26,8 +27,8 @@ class ScrapeInput(BaseModel):
|
|||
max_length=MAX_TIKTOK_SOURCES,
|
||||
description=(
|
||||
"TikTok URLs to scrape: a video, a profile (/@<user>), a hashtag "
|
||||
"(/tag/<name>), or a search URL. Provide these OR profiles/hashtags/"
|
||||
"search_queries (at least one source is required)."
|
||||
"(/tag/<name>), or a search URL. Provide these OR profiles/hashtags "
|
||||
"(at least one source is required)."
|
||||
),
|
||||
)
|
||||
profiles: list[str] = Field(
|
||||
|
|
@ -40,21 +41,11 @@ class ScrapeInput(BaseModel):
|
|||
max_length=MAX_TIKTOK_SOURCES,
|
||||
description="Hashtag names to scrape, without the leading '#'.",
|
||||
)
|
||||
search_queries: list[str] = Field(
|
||||
default_factory=list,
|
||||
max_length=MAX_TIKTOK_SOURCES,
|
||||
description=(
|
||||
"Search terms resolved via Google (site:tiktok.com) to public TikTok "
|
||||
"videos, since TikTok's own keyword search is login-walled. Slower "
|
||||
"than hashtags/urls. To find accounts by keyword, use "
|
||||
"tiktok.user_search instead."
|
||||
),
|
||||
)
|
||||
results_per_page: int = Field(
|
||||
default=10,
|
||||
ge=1,
|
||||
le=MAX_TIKTOK_ITEMS,
|
||||
description="Max videos to pull per profile/hashtag/search target.",
|
||||
description="Max videos to pull per profile/hashtag target.",
|
||||
)
|
||||
max_items: int = Field(
|
||||
default=10,
|
||||
|
|
@ -65,10 +56,9 @@ class ScrapeInput(BaseModel):
|
|||
|
||||
@model_validator(mode="after")
|
||||
def _require_a_source(self) -> ScrapeInput:
|
||||
if not any((self.urls, self.profiles, self.hashtags, self.search_queries)):
|
||||
if not any((self.urls, self.profiles, self.hashtags)):
|
||||
raise ValueError(
|
||||
"Provide at least one of 'urls', 'profiles', 'hashtags', or "
|
||||
"'search_queries'."
|
||||
"Provide at least one of 'urls', 'profiles', or 'hashtags'."
|
||||
)
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -222,9 +222,7 @@ class _RotatingSession:
|
|||
await self.close()
|
||||
self.rotations += 1
|
||||
await self._open()
|
||||
logger.info(
|
||||
"[instagram] rotated proxy session (rotation #%d)", self.rotations
|
||||
)
|
||||
logger.info("[instagram] rotated proxy session (rotation #%d)", self.rotations)
|
||||
return self.session
|
||||
|
||||
async def pace(self) -> None:
|
||||
|
|
@ -378,9 +376,7 @@ async def _fetch(
|
|||
if status == _BACKOFF_STATUS and backoffs < _MAX_BACKOFFS:
|
||||
backoffs += 1
|
||||
delay = _BACKOFF_BASE_S * (2 ** (backoffs - 1))
|
||||
logger.warning(
|
||||
"[instagram] 429 on %s; backing off %.1fs", path, delay
|
||||
)
|
||||
logger.warning("[instagram] 429 on %s; backing off %.1fs", path, delay)
|
||||
await asyncio.sleep(delay + random.uniform(0, 1))
|
||||
continue
|
||||
if status in _ROTATE_STATUSES:
|
||||
|
|
|
|||
|
|
@ -171,7 +171,9 @@ def _relay_child(node: dict[str, Any]) -> dict[str, Any]:
|
|||
mt = node.get("media_type")
|
||||
vv = node.get("video_versions")
|
||||
video_url = (
|
||||
vv[0].get("url") if isinstance(vv, list) and vv and isinstance(vv[0], dict) else None
|
||||
vv[0].get("url")
|
||||
if isinstance(vv, list) and vv and isinstance(vv[0], dict)
|
||||
else None
|
||||
)
|
||||
is_video = mt == 2 or bool(video_url)
|
||||
return {
|
||||
|
|
@ -290,9 +292,7 @@ def parse_profile(user: dict[str, Any]) -> dict[str, Any]:
|
|||
_APP_JSON_RE = re.compile(
|
||||
r'<script type="application/json"[^>]*>(.*?)</script>', re.DOTALL
|
||||
)
|
||||
_OG_RE = re.compile(
|
||||
r'<meta\s+property="og:([^"]+)"\s+content="([^"]*)"', re.IGNORECASE
|
||||
)
|
||||
_OG_RE = re.compile(r'<meta\s+property="og:([^"]+)"\s+content="([^"]*)"', re.IGNORECASE)
|
||||
# og tags are the fallback source (used only when the relay blob is absent). They
|
||||
# follow a fixed English shape because the fetch layer pins Accept-Language en-US:
|
||||
# og:description = "{likes} likes, {comments} comments - {username} on {Month D, YYYY}: "{caption}""
|
||||
|
|
@ -321,7 +321,9 @@ _MEDIA_ID_RE = re.compile(r"instagram://media\?id=(\d+)")
|
|||
def _og_date_to_iso(value: str) -> str | None:
|
||||
"""``"July 9, 2026"`` -> ``"2026-07-09"`` (date-only; og carries no time)."""
|
||||
try:
|
||||
return datetime.strptime(value, "%B %d, %Y").replace(tzinfo=UTC).date().isoformat()
|
||||
return (
|
||||
datetime.strptime(value, "%B %d, %Y").replace(tzinfo=UTC).date().isoformat()
|
||||
)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
|
@ -359,7 +361,7 @@ def _parse_og_meta(og: dict[str, str]) -> dict[str, Any]:
|
|||
elif owner_date:
|
||||
# No usable og:title: fall back to the caption after og:description's
|
||||
# date prefix — still clean (the counts/username/date are stripped).
|
||||
out["caption"] = _clean_caption(desc[owner_date.end():])
|
||||
out["caption"] = _clean_caption(desc[owner_date.end() :])
|
||||
return out
|
||||
|
||||
|
||||
|
|
@ -438,13 +440,21 @@ def _media_from_relay(
|
|||
mt = media.get("media_type")
|
||||
cap = media.get("caption")
|
||||
caption = (
|
||||
cap.get("text") if isinstance(cap, dict) else (cap if isinstance(cap, str) else None)
|
||||
cap.get("text")
|
||||
if isinstance(cap, dict)
|
||||
else (cap if isinstance(cap, str) else None)
|
||||
)
|
||||
carousel = media.get("carousel_media")
|
||||
carousel = [c for c in carousel if isinstance(c, dict)] if isinstance(carousel, list) else []
|
||||
carousel = (
|
||||
[c for c in carousel if isinstance(c, dict)]
|
||||
if isinstance(carousel, list)
|
||||
else []
|
||||
)
|
||||
vv = media.get("video_versions")
|
||||
video_url = (
|
||||
vv[0].get("url") if isinstance(vv, list) and vv and isinstance(vv[0], dict) else None
|
||||
vv[0].get("url")
|
||||
if isinstance(vv, list) and vv and isinstance(vv[0], dict)
|
||||
else None
|
||||
)
|
||||
is_video = mt == 2 or bool(video_url)
|
||||
owner = media.get("user") if isinstance(media.get("user"), dict) else {}
|
||||
|
|
@ -469,13 +479,18 @@ def _media_from_relay(
|
|||
"type": _MEDIA_TYPE.get(mt) or ("Video" if is_video else "Image"),
|
||||
"shortCode": media.get("code") or shortcode,
|
||||
"caption": caption,
|
||||
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption))) if caption else [],
|
||||
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption))) if caption else [],
|
||||
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption)))
|
||||
if caption
|
||||
else [],
|
||||
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption)))
|
||||
if caption
|
||||
else [],
|
||||
"url": url,
|
||||
"commentsCount": _int(media.get("comment_count")),
|
||||
"dimensionsHeight": _int(media.get("original_height")),
|
||||
"dimensionsWidth": _int(media.get("original_width")),
|
||||
"displayUrl": _iv2_url(media.get("image_versions2")) or media.get("display_uri"),
|
||||
"displayUrl": _iv2_url(media.get("image_versions2"))
|
||||
or media.get("display_uri"),
|
||||
"images": [
|
||||
u
|
||||
for c in carousel
|
||||
|
|
@ -535,8 +550,12 @@ def parse_post(
|
|||
"type": "Video" if is_video else "Image",
|
||||
"shortCode": shortcode,
|
||||
"caption": caption,
|
||||
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption))) if caption else [],
|
||||
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption))) if caption else [],
|
||||
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption)))
|
||||
if caption
|
||||
else [],
|
||||
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption)))
|
||||
if caption
|
||||
else [],
|
||||
"url": url,
|
||||
"commentsCount": og_meta.get("comments"),
|
||||
"displayUrl": og.get("image"),
|
||||
|
|
|
|||
|
|
@ -328,9 +328,7 @@ async def _discover_via_google(
|
|||
return resolved
|
||||
|
||||
|
||||
async def _discover(
|
||||
query: str, *, search_type: str, limit: int
|
||||
) -> list[ResolvedUrl]:
|
||||
async def _discover(query: str, *, search_type: str, limit: int) -> list[ResolvedUrl]:
|
||||
"""Resolve a discovery query into profile targets - anonymously.
|
||||
|
||||
A query that is a valid handle resolves directly against the anonymous
|
||||
|
|
@ -397,9 +395,7 @@ async def iter_instagram(
|
|||
|
||||
# posts / reels -> media feeds, de-duped by id across targets.
|
||||
jobs = [
|
||||
_media_flow(
|
||||
r, input_model=input_model, cutoff=cutoff, per_target=per_target
|
||||
)
|
||||
_media_flow(r, input_model=input_model, cutoff=cutoff, per_target=per_target)
|
||||
for r in targets
|
||||
]
|
||||
seen: set[str] = set()
|
||||
|
|
|
|||
|
|
@ -25,9 +25,7 @@ from urllib.parse import urlparse
|
|||
|
||||
ResolvedKind = Literal["profile", "post", "reel"]
|
||||
|
||||
_INSTAGRAM_HOSTS = frozenset(
|
||||
{"m.instagram.com", "www.instagram.com", "instagram.com"}
|
||||
)
|
||||
_INSTAGRAM_HOSTS = frozenset({"m.instagram.com", "www.instagram.com", "instagram.com"})
|
||||
_STRIP_SEGMENTS = frozenset({"_u", "profilecard"})
|
||||
_RESERVED = frozenset(
|
||||
{"p", "s", "tv", "reel", "reels", "share", "explore", "stories", "accounts"}
|
||||
|
|
@ -68,9 +66,7 @@ def resolve_url(url: str) -> ResolvedUrl | None:
|
|||
if "instagram.com" not in url.lower():
|
||||
token = url.strip().lstrip("@")
|
||||
if token and "/" not in token and "." not in token:
|
||||
return ResolvedUrl(
|
||||
"profile", token, f"https://www.instagram.com/{token}/"
|
||||
)
|
||||
return ResolvedUrl("profile", token, f"https://www.instagram.com/{token}/")
|
||||
segments = _segments(url)
|
||||
if not segments:
|
||||
return None
|
||||
|
|
@ -83,9 +79,7 @@ def resolve_url(url: str) -> ResolvedUrl | None:
|
|||
return ResolvedUrl("reel", code, url, numeric_post_id=code.isdigit())
|
||||
if head == "stories" and len(segments) >= 2:
|
||||
user = segments[1]
|
||||
return ResolvedUrl(
|
||||
"profile", user, f"https://www.instagram.com/{user}/"
|
||||
)
|
||||
return ResolvedUrl("profile", user, f"https://www.instagram.com/{user}/")
|
||||
if head not in _RESERVED:
|
||||
return ResolvedUrl("profile", head, url)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -5,11 +5,11 @@ from __future__ import annotations
|
|||
from datetime import UTC, datetime
|
||||
|
||||
|
||||
def epoch_to_iso(seconds: int | None) -> str | None:
|
||||
def epoch_to_iso(seconds: int | str | None) -> str | None:
|
||||
"""Convert a Unix-seconds timestamp to ``YYYY-MM-DDTHH:MM:SS.000Z``."""
|
||||
if not seconds:
|
||||
return None
|
||||
stamp = datetime.fromtimestamp(seconds, tz=UTC)
|
||||
stamp = datetime.fromtimestamp(int(seconds), tz=UTC)
|
||||
return stamp.strftime("%Y-%m-%dT%H:%M:%S.000Z")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,9 @@ from ..targets.types import TikTokTarget
|
|||
from . import FetchFn
|
||||
|
||||
|
||||
async def iter_video(target: TikTokTarget, *, fetch: FetchFn) -> AsyncIterator[dict[str, Any]]:
|
||||
async def iter_video(
|
||||
target: TikTokTarget, *, fetch: FetchFn
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
html = await fetch(target.url)
|
||||
if not html:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -10,10 +10,6 @@ from __future__ import annotations
|
|||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from urllib.parse import quote
|
||||
|
||||
from app.proprietary.platforms.google_search.schemas import GoogleSearchScrapeInput
|
||||
from app.proprietary.platforms.google_search.scraper import scrape_serps
|
||||
|
||||
from .extraction.timestamps import now_iso
|
||||
from .flows import FetchCommentsFn, FetchFn, FetchListingFn, FetchUsersFn
|
||||
|
|
@ -35,26 +31,15 @@ from .targets.types import TikTokTarget
|
|||
|
||||
_PROFILE_URL = "https://www.tiktok.com/@{name}"
|
||||
_HASHTAG_URL = "https://www.tiktok.com/tag/{tag}"
|
||||
_SEARCH_URL = "https://www.tiktok.com/search?q={query}"
|
||||
_EXPLORE_URL = "https://www.tiktok.com/explore"
|
||||
|
||||
# A ``searchQueries`` term whose Google discovery surfaced no scrapable video
|
||||
# URLs degrades to one honest ErrorItem (mirrors the listing flow's contract:
|
||||
# never vanish silently).
|
||||
_EMPTY_DISCOVERY_MESSAGE = (
|
||||
"No public TikTok videos found for this query via Google discovery. Try a "
|
||||
"narrower phrasing, a hashtag, or a direct video URL."
|
||||
)
|
||||
|
||||
|
||||
def _resolve_targets(input_model: TikTokScrapeInput) -> list[TikTokTarget]:
|
||||
"""Build the target list from the URL/profile/hashtag sources.
|
||||
|
||||
``searchQueries`` is deliberately excluded: TikTok's own keyword search is
|
||||
login-walled for anonymous sessions, so it is routed through Google video
|
||||
discovery in :func:`iter_tiktok` instead. A raw ``tiktok.com/search?...``
|
||||
URL passed explicitly in ``startUrls``/``postURLs`` still resolves here and
|
||||
keeps its native listing routing.
|
||||
A raw ``tiktok.com/search?...`` URL passed explicitly in
|
||||
``startUrls``/``postURLs`` still resolves here and keeps its native listing
|
||||
routing; there is no keyword-search shortcut.
|
||||
"""
|
||||
targets: list[TikTokTarget] = []
|
||||
for entry in input_model.startUrls:
|
||||
|
|
@ -73,39 +58,6 @@ def _resolve_targets(input_model: TikTokScrapeInput) -> list[TikTokTarget]:
|
|||
return targets
|
||||
|
||||
|
||||
async def _discover_via_google(query: str, *, limit: int) -> list[TikTokTarget]:
|
||||
"""Discover public TikTok video targets via Google ``site:tiktok.com``.
|
||||
|
||||
TikTok's anonymous keyword search is login-walled, so we reuse the existing
|
||||
``google_search`` platform, classify each organic URL with ``resolve_target``,
|
||||
and keep only video hits (``/@user/video/<id>``) — the one kind that scrapes
|
||||
reliably over plain HTTP. Profile/hashtag/search/photo/non-tiktok results are
|
||||
dropped (accounts belong to the ``user_search`` verb). De-duped, capped at
|
||||
``limit``.
|
||||
"""
|
||||
serps = await scrape_serps(
|
||||
GoogleSearchScrapeInput(
|
||||
queries=query, site="tiktok.com", maxPagesPerQuery=1
|
||||
),
|
||||
limit=1,
|
||||
)
|
||||
resolved: list[TikTokTarget] = []
|
||||
seen: set[str] = set()
|
||||
for serp in serps:
|
||||
for org in serp.get("organicResults") or []:
|
||||
url = org.get("url", "") if isinstance(org, dict) else ""
|
||||
target = resolve_target(url)
|
||||
if target is None or target.kind != "video":
|
||||
continue
|
||||
if target.value in seen:
|
||||
continue
|
||||
seen.add(target.value)
|
||||
resolved.append(target)
|
||||
if len(resolved) >= limit:
|
||||
return resolved
|
||||
return resolved
|
||||
|
||||
|
||||
def _dispatch(
|
||||
target: TikTokTarget,
|
||||
*,
|
||||
|
|
@ -128,11 +80,9 @@ async def iter_tiktok(
|
|||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""Yield normalized items for every resolved target, in order.
|
||||
|
||||
Direct sources (URLs, profiles, hashtags) resolve up front; ``searchQueries``
|
||||
then run through Google video discovery. The video flow's ``fetch_html``
|
||||
opens its own warmed proxy session per call when none is bound; the listing
|
||||
flow drives its own browser. Neither binds a ContextVar across these
|
||||
``yield``s, so the generator stays context-safe.
|
||||
The video flow's ``fetch_html`` opens its own warmed proxy session per call
|
||||
when none is bound; the listing flow drives its own browser. Neither binds a
|
||||
ContextVar across these ``yield``s, so the generator stays context-safe.
|
||||
"""
|
||||
cap = input_model.resultsPerPage
|
||||
for target in _resolve_targets(input_model):
|
||||
|
|
@ -141,27 +91,6 @@ async def iter_tiktok(
|
|||
):
|
||||
yield item
|
||||
|
||||
# searchQueries -> Google-discovered public video URLs, de-duped across
|
||||
# queries so the same video surfacing under two terms is scraped once.
|
||||
seen_videos: set[str] = set()
|
||||
for query in input_model.searchQueries:
|
||||
discovered = await _discover_via_google(query, limit=cap)
|
||||
if not discovered:
|
||||
yield ErrorItem(
|
||||
url=_SEARCH_URL.format(query=quote(query)),
|
||||
input=query,
|
||||
error=_EMPTY_DISCOVERY_MESSAGE,
|
||||
errorCode="no_items",
|
||||
scrapedAt=now_iso(),
|
||||
).to_output()
|
||||
continue
|
||||
for target in discovered:
|
||||
if target.value in seen_videos:
|
||||
continue
|
||||
seen_videos.add(target.value)
|
||||
async for item in iter_video(target, fetch=fetch):
|
||||
yield item
|
||||
|
||||
|
||||
async def scrape_tiktok(
|
||||
input_model: TikTokScrapeInput,
|
||||
|
|
@ -174,7 +103,9 @@ async def scrape_tiktok(
|
|||
from app.capabilities.core.progress import emit_progress
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
async for item in iter_tiktok(input_model, fetch=fetch, fetch_listing=fetch_listing):
|
||||
async for item in iter_tiktok(
|
||||
input_model, fetch=fetch, fetch_listing=fetch_listing
|
||||
):
|
||||
results.append(item)
|
||||
emit_progress("scraping", current=len(results), total=limit, unit="item")
|
||||
if limit is not None and len(results) >= limit:
|
||||
|
|
|
|||
|
|
@ -54,9 +54,7 @@ from app.proprietary.platforms.instagram.url_resolver import resolve_url # noqa
|
|||
_PROFILE = "natgeo"
|
||||
_SEARCH_TERM = "national geographic"
|
||||
|
||||
_FIXTURE_DIR = (
|
||||
_BACKEND_ROOT / "tests" / "unit" / "platforms" / "instagram" / "fixtures"
|
||||
)
|
||||
_FIXTURE_DIR = _BACKEND_ROOT / "tests" / "unit" / "platforms" / "instagram" / "fixtures"
|
||||
|
||||
# Fields to strip from dumped fixtures so we never commit PII / volatile tokens.
|
||||
_PII_KEYS = frozenset(
|
||||
|
|
@ -98,7 +96,9 @@ async def step0_probe() -> bool:
|
|||
data = await fetch_json(
|
||||
"api/v1/users/web_profile_info/", {"username": _PROFILE}
|
||||
)
|
||||
user = (data or {}).get("data", {}).get("user") if isinstance(data, dict) else None
|
||||
user = (
|
||||
(data or {}).get("data", {}).get("user") if isinstance(data, dict) else None
|
||||
)
|
||||
print(f" web_profile_info({_PROFILE}) -> user={'yes' if user else 'no'}")
|
||||
return _check("sticky web_profile_info", minted and bool(user))
|
||||
|
||||
|
|
@ -179,9 +179,7 @@ async def step5_search() -> bool:
|
|||
|
||||
async def step6_dump_fixtures(post_url: str | None) -> bool:
|
||||
_hr("STEP 6 — dump trimmed, anonymized fixtures for offline tests")
|
||||
profile = await fetch_json(
|
||||
"api/v1/users/web_profile_info/", {"username": _PROFILE}
|
||||
)
|
||||
profile = await fetch_json("api/v1/users/web_profile_info/", {"username": _PROFILE})
|
||||
_FIXTURE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
wrote = []
|
||||
if isinstance(profile, dict) and profile.get("data", {}).get("user"):
|
||||
|
|
|
|||
|
|
@ -178,7 +178,9 @@ async def stage_pipeline() -> bool:
|
|||
f"{len(items)} item(s)",
|
||||
)
|
||||
if items:
|
||||
print(f" sample: {items[0].get('webVideoUrl')} — {items[0].get('text', '')[:60]!r}")
|
||||
print(
|
||||
f" sample: {items[0].get('webVideoUrl')} — {items[0].get('text', '')[:60]!r}"
|
||||
)
|
||||
return ok
|
||||
|
||||
|
||||
|
|
@ -210,9 +212,7 @@ async def stage_comments(video_url: str) -> tuple[bool, list[dict[str, Any]]]:
|
|||
# Comments load over a signed /api/comment/list XHR that TikTok serves to
|
||||
# anonymous sessions once the panel opens. Pass if real comments come back
|
||||
# OR a graceful ErrorItem (video has none / disabled / withheld).
|
||||
items = await scrape_tiktok_comments(
|
||||
[video_url], per_video=_COUNT, limit=_COUNT
|
||||
)
|
||||
items = await scrape_tiktok_comments([video_url], per_video=_COUNT, limit=_COUNT)
|
||||
has_comment = any(it.get("id") and not it.get("errorCode") for it in items)
|
||||
has_error = any(it.get("errorCode") == "no_comments" for it in items)
|
||||
ok = _check(
|
||||
|
|
@ -253,7 +253,9 @@ async def stage_trending() -> tuple[bool, list[dict[str, Any]]]:
|
|||
f"{len(items)} item(s); videos={len(real)}",
|
||||
)
|
||||
if real:
|
||||
print(f" sample: {real[0].get('webVideoUrl')} — {real[0].get('text', '')[:60]!r}")
|
||||
print(
|
||||
f" sample: {real[0].get('webVideoUrl')} — {real[0].get('text', '')[:60]!r}"
|
||||
)
|
||||
return ok, items
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -66,6 +66,4 @@ def test_details_wraps_profile_items():
|
|||
|
||||
def test_details_rejects_both_sources():
|
||||
with pytest.raises(ValidationError):
|
||||
DetailsInput(
|
||||
urls=["https://www.instagram.com/natgeo/"], search_queries=["x"]
|
||||
)
|
||||
DetailsInput(urls=["https://www.instagram.com/natgeo/"], search_queries=["x"])
|
||||
|
|
|
|||
|
|
@ -54,7 +54,6 @@ async def test_forwards_typed_sources_and_limit():
|
|||
ScrapeInput(
|
||||
profiles=["nasa"],
|
||||
hashtags=["food"],
|
||||
search_queries=["cats"],
|
||||
results_per_page=7,
|
||||
max_items=25,
|
||||
)
|
||||
|
|
@ -63,7 +62,6 @@ async def test_forwards_typed_sources_and_limit():
|
|||
(actor_input, limit) = scraper.calls[0]
|
||||
assert actor_input.profiles == ["nasa"]
|
||||
assert actor_input.hashtags == ["food"]
|
||||
assert actor_input.searchQueries == ["cats"]
|
||||
assert actor_input.resultsPerPage == 7
|
||||
# The outer collection limit is the caller's total-item cap.
|
||||
assert limit == 25
|
||||
|
|
|
|||
|
|
@ -42,8 +42,7 @@ def _profile_payload(n: int) -> dict:
|
|||
"edge_owner_to_timeline_media": {
|
||||
"count": n,
|
||||
"edges": [
|
||||
{"node": {"id": str(i), "shortcode": f"S{i}"}}
|
||||
for i in range(n)
|
||||
{"node": {"id": str(i), "shortcode": f"S{i}"}} for i in range(n)
|
||||
],
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -33,9 +33,7 @@ async def test_google_discovery_keeps_only_profiles(monkeypatch):
|
|||
"https://example.com/not-instagram",
|
||||
),
|
||||
)
|
||||
targets = await scraper._discover(
|
||||
"nat geo photos", search_type="profile", limit=10
|
||||
)
|
||||
targets = await scraper._discover("nat geo photos", search_type="profile", limit=10)
|
||||
assert [(t.kind, t.value) for t in targets] == [("profile", "natgeo")]
|
||||
|
||||
|
||||
|
|
@ -48,9 +46,7 @@ async def test_google_discovery_dedupes(monkeypatch):
|
|||
"https://www.instagram.com/natgeo/",
|
||||
),
|
||||
)
|
||||
targets = await scraper._discover(
|
||||
"nat geo photos", search_type="profile", limit=10
|
||||
)
|
||||
targets = await scraper._discover("nat geo photos", search_type="profile", limit=10)
|
||||
assert len(targets) == 1
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -102,7 +102,9 @@ async def test_warms_then_returns_json():
|
|||
holder = _FakeHolder([_FakeSession(200, csrftoken=True)])
|
||||
token = _current_session.set(holder)
|
||||
try:
|
||||
result = await fetch_json("api/v1/users/web_profile_info/", {"username": "natgeo"})
|
||||
result = await fetch_json(
|
||||
"api/v1/users/web_profile_info/", {"username": "natgeo"}
|
||||
)
|
||||
finally:
|
||||
_current_session.reset(token)
|
||||
assert result == _PAYLOAD
|
||||
|
|
|
|||
|
|
@ -175,8 +175,14 @@ def test_parse_post_prefers_relay_json():
|
|||
"image_versions2": {"candidates": [{"url": "https://cdn/c2.jpg"}]},
|
||||
},
|
||||
],
|
||||
"usertags": {"in": [{"position": [0.5, 0.5], "user": {"username": "tagged1", "id": "77"}}]},
|
||||
"coauthor_producers": [{"username": "coauthor1", "id": "88", "is_verified": True}],
|
||||
"usertags": {
|
||||
"in": [
|
||||
{"position": [0.5, 0.5], "user": {"username": "tagged1", "id": "77"}}
|
||||
]
|
||||
},
|
||||
"coauthor_producers": [
|
||||
{"username": "coauthor1", "id": "88", "is_verified": True}
|
||||
],
|
||||
"location": {"id": "123", "name": "Bali"},
|
||||
}
|
||||
html = (
|
||||
|
|
|
|||
|
|
@ -1,149 +0,0 @@
|
|||
"""Offline tests for Google-backed TikTok video discovery.
|
||||
|
||||
``searchQueries`` are login-walled on TikTok's native search, so they route
|
||||
through the ``google_search`` platform (``site:tiktok.com``): each organic URL
|
||||
is classified with ``resolve_target`` and only video hits (``/@user/video/<id>``)
|
||||
are kept — profiles/hashtags/search/photo/non-tiktok are dropped (accounts
|
||||
belong to the user-search verb). These tests inject a fake ``scrape_serps`` so
|
||||
there is no network: they pin the classification, cross-query de-dup, the limit
|
||||
cap, the barren-query ErrorItem, and that no ``/search?q=`` listing target is
|
||||
ever built.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from app.proprietary.platforms.tiktok import (
|
||||
TikTokScrapeInput,
|
||||
orchestrator,
|
||||
scrape_tiktok,
|
||||
)
|
||||
|
||||
|
||||
def _fake_serps(*organic_urls: str):
|
||||
async def _scrape_serps(input_model, *, limit=None):
|
||||
assert input_model.site == "tiktok.com"
|
||||
assert input_model.maxPagesPerQuery == 1
|
||||
return [{"organicResults": [{"url": u} for u in organic_urls]}]
|
||||
|
||||
return _scrape_serps
|
||||
|
||||
|
||||
def _video_page(url: str) -> str:
|
||||
"""Render a rehydration blob for a ``/@user/video/<id>`` URL."""
|
||||
video_id = url.rsplit("/", 1)[1]
|
||||
username = url.split("@")[1].split("/")[0]
|
||||
blob = {
|
||||
"__DEFAULT_SCOPE__": {
|
||||
"webapp.video-detail": {
|
||||
"itemInfo": {
|
||||
"itemStruct": {
|
||||
"id": video_id,
|
||||
"desc": "hi",
|
||||
"author": {"uniqueId": username},
|
||||
"stats": {"diggCount": 1},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return (
|
||||
'<script id="__UNIVERSAL_DATA_FOR_REHYDRATION__" '
|
||||
f'type="application/json">{json.dumps(blob)}</script>'
|
||||
)
|
||||
|
||||
|
||||
async def _fetch_video(url: str) -> str:
|
||||
return _video_page(url)
|
||||
|
||||
|
||||
async def test_search_discovery_keeps_only_videos(monkeypatch):
|
||||
# Only the video URL survives; profile / hashtag / search / photo /
|
||||
# non-tiktok organic results are dropped.
|
||||
monkeypatch.setattr(
|
||||
orchestrator,
|
||||
"scrape_serps",
|
||||
_fake_serps(
|
||||
"https://www.tiktok.com/@nasa/video/123",
|
||||
"https://www.tiktok.com/@nasa",
|
||||
"https://www.tiktok.com/tag/space",
|
||||
"https://www.tiktok.com/search?q=space",
|
||||
"https://www.tiktok.com/@nasa/photo/999",
|
||||
"https://example.com/not-tiktok",
|
||||
),
|
||||
)
|
||||
items = await scrape_tiktok(
|
||||
TikTokScrapeInput(searchQueries=["space"], resultsPerPage=10),
|
||||
fetch=_fetch_video,
|
||||
)
|
||||
assert [i["id"] for i in items] == ["123"]
|
||||
|
||||
|
||||
async def test_search_discovery_dedupes_across_queries(monkeypatch):
|
||||
# The same video surfacing under two queries is scraped once.
|
||||
monkeypatch.setattr(
|
||||
orchestrator,
|
||||
"scrape_serps",
|
||||
_fake_serps("https://www.tiktok.com/@nasa/video/123"),
|
||||
)
|
||||
items = await scrape_tiktok(
|
||||
TikTokScrapeInput(searchQueries=["space", "rockets"], resultsPerPage=10),
|
||||
fetch=_fetch_video,
|
||||
)
|
||||
assert [i["id"] for i in items] == ["123"]
|
||||
|
||||
|
||||
async def test_search_discovery_respects_per_target_limit(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
orchestrator,
|
||||
"scrape_serps",
|
||||
_fake_serps(
|
||||
"https://www.tiktok.com/@a/video/1",
|
||||
"https://www.tiktok.com/@b/video/2",
|
||||
"https://www.tiktok.com/@c/video/3",
|
||||
),
|
||||
)
|
||||
items = await scrape_tiktok(
|
||||
TikTokScrapeInput(searchQueries=["x"], resultsPerPage=2),
|
||||
fetch=_fetch_video,
|
||||
)
|
||||
assert [i["id"] for i in items] == ["1", "2"]
|
||||
|
||||
|
||||
async def test_search_barren_query_emits_error_item(monkeypatch):
|
||||
# A query whose discovery finds no video URLs degrades to one ErrorItem.
|
||||
monkeypatch.setattr(
|
||||
orchestrator,
|
||||
"scrape_serps",
|
||||
_fake_serps(
|
||||
"https://www.tiktok.com/@nasa",
|
||||
"https://example.com/x",
|
||||
),
|
||||
)
|
||||
items = await scrape_tiktok(
|
||||
TikTokScrapeInput(searchQueries=["space"], resultsPerPage=10),
|
||||
fetch=_fetch_video,
|
||||
)
|
||||
assert len(items) == 1
|
||||
assert items[0]["errorCode"] == "no_items"
|
||||
assert items[0]["input"] == "space"
|
||||
|
||||
|
||||
async def test_search_never_builds_listing_target(monkeypatch):
|
||||
# searchQueries must never hit the (login-walled) native search listing flow.
|
||||
monkeypatch.setattr(
|
||||
orchestrator,
|
||||
"scrape_serps",
|
||||
_fake_serps("https://www.tiktok.com/@nasa/video/123"),
|
||||
)
|
||||
|
||||
async def _boom_listing(_url: str, _count: int) -> list[dict]:
|
||||
raise AssertionError("searchQueries must not build a listing target")
|
||||
|
||||
items = await scrape_tiktok(
|
||||
TikTokScrapeInput(searchQueries=["space"], resultsPerPage=10),
|
||||
fetch=_fetch_video,
|
||||
fetch_listing=_boom_listing,
|
||||
)
|
||||
assert [i["id"] for i in items] == ["123"]
|
||||
|
|
@ -84,7 +84,9 @@ async def test_warms_then_returns_html():
|
|||
|
||||
|
||||
async def test_rotates_when_warm_fails_then_succeeds():
|
||||
holder = _FakeHolder([_FakeSession(200, warms=False), _FakeSession(200, warms=True)])
|
||||
holder = _FakeHolder(
|
||||
[_FakeSession(200, warms=False), _FakeSession(200, warms=True)]
|
||||
)
|
||||
token = _current_session.set(holder)
|
||||
try:
|
||||
result = await client.fetch_html("https://www.tiktok.com/@scout2015")
|
||||
|
|
@ -119,9 +121,7 @@ async def test_rotates_and_rewarms_on_403():
|
|||
|
||||
async def test_persistent_403_raises_blocked(monkeypatch):
|
||||
_no_sleep(monkeypatch)
|
||||
holder = _FakeHolder(
|
||||
[_FakeSession(403) for _ in range(client._MAX_ROTATIONS + 1)]
|
||||
)
|
||||
holder = _FakeHolder([_FakeSession(403) for _ in range(client._MAX_ROTATIONS + 1)])
|
||||
token = _current_session.set(holder)
|
||||
try:
|
||||
raised = False
|
||||
|
|
|
|||
|
|
@ -6,7 +6,9 @@ from app.proprietary.platforms.tiktok.targets import resolve_target
|
|||
|
||||
|
||||
def test_resolve_video_carries_username_and_id():
|
||||
target = resolve_target("https://www.tiktok.com/@scout2015/video/6718335390845095173")
|
||||
target = resolve_target(
|
||||
"https://www.tiktok.com/@scout2015/video/6718335390845095173"
|
||||
)
|
||||
assert target is not None
|
||||
assert target.kind == "video"
|
||||
assert target.value == "6718335390845095173"
|
||||
|
|
|
|||
|
|
@ -28,9 +28,7 @@ async def test_user_search_parses_dedupes_and_caps():
|
|||
async def fake_fetch(_url: str, _cap: int) -> list[dict]:
|
||||
return [_user("1", "nasa"), _user("1", "nasa"), _user("2", "nasa2")]
|
||||
|
||||
items = await search_tiktok_users(
|
||||
["nasa"], per_query=2, fetch_users=fake_fetch
|
||||
)
|
||||
items = await search_tiktok_users(["nasa"], per_query=2, fetch_users=fake_fetch)
|
||||
|
||||
assert [i["id"] for i in items] == ["1", "2"]
|
||||
first = items[0]
|
||||
|
|
@ -49,9 +47,7 @@ async def test_user_search_empty_query_emits_error_item():
|
|||
async def fake_fetch(_url: str, _cap: int) -> list[dict]:
|
||||
return []
|
||||
|
||||
items = await search_tiktok_users(
|
||||
["ghost"], per_query=5, fetch_users=fake_fetch
|
||||
)
|
||||
items = await search_tiktok_users(["ghost"], per_query=5, fetch_users=fake_fetch)
|
||||
|
||||
assert len(items) == 1
|
||||
assert items[0]["errorCode"] == "no_users"
|
||||
|
|
|
|||
|
|
@ -58,9 +58,7 @@ def test_chat_model_requires_slash_tools_and_context():
|
|||
|
||||
|
||||
def test_excluded_provider_slug_is_filtered():
|
||||
assert not is_requesty_chat_model(
|
||||
_requesty_model(model_id="amazon/nova-pro-v1")
|
||||
)
|
||||
assert not is_requesty_chat_model(_requesty_model(model_id="amazon/nova-pro-v1"))
|
||||
|
||||
|
||||
def test_image_generation_models_excluded_from_chat_and_flagged():
|
||||
|
|
@ -89,9 +87,7 @@ def test_normalize_maps_context_window_and_capabilities():
|
|||
name="GPT-4o mini",
|
||||
),
|
||||
_requesty_model(model_id="openai/gpt-4o-mini", tools=False),
|
||||
_requesty_model(
|
||||
model_id="black-forest-labs/flux", image_generation=True
|
||||
),
|
||||
_requesty_model(model_id="black-forest-labs/flux", image_generation=True),
|
||||
]
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -21,8 +21,7 @@ PDFS = REPO / "data" / "multimodal_doc" / "mmlongbench" / "pdfs"
|
|||
|
||||
def main() -> None:
|
||||
rows = [
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip()
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
|
||||
]
|
||||
|
||||
# 1) SSL clustering: failures by question index per arm
|
||||
|
|
@ -35,11 +34,19 @@ def main() -> None:
|
|||
arm_seen_count[arm] += 1
|
||||
qid_order[f"{arm}::{row['qid']}"] = idx
|
||||
err = row.get("error") or ""
|
||||
cluster = "ssl" if "SSLError" in err else (
|
||||
"empty" if not (row.get("raw_text") or "").strip() and not err else (
|
||||
"5xx" if "502" in err or "503" in err else (
|
||||
"size_limit" if "exceeds" in err.lower() and "limit" in err.lower() else (
|
||||
"other_err" if err else "ok"
|
||||
cluster = (
|
||||
"ssl"
|
||||
if "SSLError" in err
|
||||
else (
|
||||
"empty"
|
||||
if not (row.get("raw_text") or "").strip() and not err
|
||||
else (
|
||||
"5xx"
|
||||
if "502" in err or "503" in err
|
||||
else (
|
||||
"size_limit"
|
||||
if "exceeds" in err.lower() and "limit" in err.lower()
|
||||
else ("other_err" if err else "ok")
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -100,19 +107,26 @@ def main() -> None:
|
|||
err = row.get("error") or ""
|
||||
empty = not (row.get("raw_text") or "").strip()
|
||||
if err or empty:
|
||||
by_pdf[row["doc_id"]].append({
|
||||
"arm": row["arm"],
|
||||
"qid": row["qid"],
|
||||
"err_kind": (
|
||||
"ssl" if "SSLError" in err
|
||||
else "size_limit" if "exceeds" in err.lower() and "limit" in err.lower()
|
||||
else "5xx" if "502" in err or "503" in err
|
||||
else "json_decode" if "JSONDecodeError" in err
|
||||
else "empty" if empty and not err
|
||||
else "other"
|
||||
),
|
||||
"pages": row.get("pages"),
|
||||
})
|
||||
by_pdf[row["doc_id"]].append(
|
||||
{
|
||||
"arm": row["arm"],
|
||||
"qid": row["qid"],
|
||||
"err_kind": (
|
||||
"ssl"
|
||||
if "SSLError" in err
|
||||
else "size_limit"
|
||||
if "exceeds" in err.lower() and "limit" in err.lower()
|
||||
else "5xx"
|
||||
if "502" in err or "503" in err
|
||||
else "json_decode"
|
||||
if "JSONDecodeError" in err
|
||||
else "empty"
|
||||
if empty and not err
|
||||
else "other"
|
||||
),
|
||||
"pages": row.get("pages"),
|
||||
}
|
||||
)
|
||||
for doc, items in sorted(by_pdf.items(), key=lambda x: (-len(x[1]), x[0])):
|
||||
kinds = Counter(i["err_kind"] for i in items)
|
||||
arms = sorted({i["arm"] for i in items})
|
||||
|
|
|
|||
|
|
@ -51,8 +51,7 @@ def _classify(error: str | None, raw_text: str) -> str:
|
|||
|
||||
def main() -> None:
|
||||
rows = [
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip()
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
|
||||
]
|
||||
|
||||
by_arm_failures: dict[str, list[dict]] = defaultdict(list)
|
||||
|
|
@ -121,7 +120,9 @@ def main() -> None:
|
|||
print("=" * 90)
|
||||
for entry in by_arm_failures.get("native_pdf", []):
|
||||
err = (entry["error"] or "(no error string)")[:240].replace("\n", " ")
|
||||
print(f" {entry['qid']} doc={entry['doc_id']} pages={entry['pages']} cluster={entry['cluster']}")
|
||||
print(
|
||||
f" {entry['qid']} doc={entry['doc_id']} pages={entry['pages']} cluster={entry['cluster']}"
|
||||
)
|
||||
print(f" err: {err}")
|
||||
|
||||
summary: dict[str, Any] = {
|
||||
|
|
@ -130,18 +131,13 @@ def main() -> None:
|
|||
"n": n_per_arm[arm],
|
||||
"failures": len(by_arm_failures[arm]),
|
||||
"rate": len(by_arm_failures[arm]) / n_per_arm[arm],
|
||||
"clusters": {
|
||||
cluster: len(items)
|
||||
for cluster, items in error_clusters[arm].items()
|
||||
},
|
||||
"clusters": {cluster: len(items) for cluster, items in error_clusters[arm].items()},
|
||||
"rows": by_arm_failures[arm],
|
||||
}
|
||||
for arm in sorted(n_per_arm)
|
||||
},
|
||||
"per_pdf": {
|
||||
pdf: [
|
||||
{**r, "arm": r["arm"]} for r in failures
|
||||
]
|
||||
pdf: [{**r, "arm": r["arm"]} for r in failures]
|
||||
for pdf, failures in by_pdf_failures.items()
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,9 +23,7 @@ SAFE_CHARS = (CTX_TOKENS - PROMPT_OVERHEAD_TOKENS - MAX_OUTPUT_TOKENS) * CHARS_P
|
|||
|
||||
def main() -> None:
|
||||
rows = [
|
||||
json.loads(line)
|
||||
for line in MAP.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip()
|
||||
json.loads(line) for line in MAP.read_text(encoding="utf-8").splitlines() if line.strip()
|
||||
]
|
||||
|
||||
total = len(rows)
|
||||
|
|
|
|||
|
|
@ -67,14 +67,18 @@ def classify(error: str | None, raw_text: str) -> str:
|
|||
|
||||
def main() -> None:
|
||||
rows = [
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip()
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
|
||||
]
|
||||
by_arm: dict[str, dict] = defaultdict(lambda: {
|
||||
"n": 0, "correct": 0,
|
||||
"transient_ssl_or_5xx": 0, "transient_empty": 0,
|
||||
"intrinsic_limit": 0, "other_error": 0,
|
||||
})
|
||||
by_arm: dict[str, dict] = defaultdict(
|
||||
lambda: {
|
||||
"n": 0,
|
||||
"correct": 0,
|
||||
"transient_ssl_or_5xx": 0,
|
||||
"transient_empty": 0,
|
||||
"intrinsic_limit": 0,
|
||||
"other_error": 0,
|
||||
}
|
||||
)
|
||||
for row in rows:
|
||||
arm = row["arm"]
|
||||
m = by_arm[arm]
|
||||
|
|
@ -86,7 +90,9 @@ def main() -> None:
|
|||
if kind != "ok":
|
||||
m[kind] += 1
|
||||
|
||||
print(f"{'arm':<25} {'raw acc%':>8} {'transient':>10} {'intrinsic':>10} {'other':>6} {'adj acc% (no transient)':>22}")
|
||||
print(
|
||||
f"{'arm':<25} {'raw acc%':>8} {'transient':>10} {'intrinsic':>10} {'other':>6} {'adj acc% (no transient)':>22}"
|
||||
)
|
||||
print("-" * 88)
|
||||
for arm in sorted(by_arm):
|
||||
m = by_arm[arm]
|
||||
|
|
@ -96,9 +102,7 @@ def main() -> None:
|
|||
other = m["other_error"]
|
||||
usable = m["n"] - transient
|
||||
adj = m["correct"] / usable * 100 if usable else 0
|
||||
print(
|
||||
f"{arm:<25} {raw:>7.1f}% {transient:>10} {intrinsic:>10} {other:>6} {adj:>21.1f}%"
|
||||
)
|
||||
print(f"{arm:<25} {raw:>7.1f}% {transient:>10} {intrinsic:>10} {other:>6} {adj:>21.1f}%")
|
||||
|
||||
print()
|
||||
print("transient = SSLError / 502 / 503 / empty stream / mid-stream JSON decode (would")
|
||||
|
|
|
|||
|
|
@ -95,7 +95,7 @@ def _mcnemar_exact_pvalue(b: int, c: int) -> float:
|
|||
k = min(b, c)
|
||||
# Two-sided exact: 2 * P(X <= k) clipped at 1.0
|
||||
cdf = sum(_binom_coef(n, i) for i in range(k + 1))
|
||||
p = 2.0 * cdf / (2 ** n)
|
||||
p = 2.0 * cdf / (2**n)
|
||||
return min(1.0, p)
|
||||
|
||||
|
||||
|
|
@ -116,7 +116,7 @@ def _mcnemar_table(rows: list[dict]) -> dict:
|
|||
qids = sorted(by_qid)
|
||||
out: dict[str, dict] = {"arms": arms, "n_qids": len(qids), "pairs": []}
|
||||
for i, ai in enumerate(arms):
|
||||
for aj in arms[i + 1:]:
|
||||
for aj in arms[i + 1 :]:
|
||||
b = c = both = neither = 0
|
||||
for q in qids:
|
||||
row = by_qid[q]
|
||||
|
|
@ -132,12 +132,17 @@ def _mcnemar_table(rows: list[dict]) -> dict:
|
|||
else:
|
||||
neither += 1
|
||||
p = _mcnemar_exact_pvalue(b, c)
|
||||
out["pairs"].append({
|
||||
"arm_i": ai, "arm_j": aj,
|
||||
"b_i_only": b, "c_j_only": c,
|
||||
"both_correct": both, "both_wrong": neither,
|
||||
"p_value": p,
|
||||
})
|
||||
out["pairs"].append(
|
||||
{
|
||||
"arm_i": ai,
|
||||
"arm_j": aj,
|
||||
"b_i_only": b,
|
||||
"c_j_only": c,
|
||||
"both_correct": both,
|
||||
"both_wrong": neither,
|
||||
"p_value": p,
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
|
|
@ -154,9 +159,7 @@ def _per_pdf_stats(rows: list[dict]) -> dict[str, dict]:
|
|||
arm = r["arm"]
|
||||
pdf = r["doc_id"]
|
||||
graded = r.get("graded") or {}
|
||||
bucket.setdefault(arm, {}).setdefault(pdf, []).append(
|
||||
bool(graded.get("correct"))
|
||||
)
|
||||
bucket.setdefault(arm, {}).setdefault(pdf, []).append(bool(graded.get("correct")))
|
||||
|
||||
out: dict[str, dict] = {}
|
||||
for arm, pdfs in bucket.items():
|
||||
|
|
@ -207,7 +210,8 @@ def _per_arm_latency(rows: list[dict]) -> dict[str, dict]:
|
|||
# Coefficient of variation: std / mean (unitless tail-fatness).
|
||||
"cv": (
|
||||
statistics.stdev(lats) / statistics.mean(lats)
|
||||
if len(lats) > 1 and statistics.mean(lats) > 0 else 0.0
|
||||
if len(lats) > 1 and statistics.mean(lats) > 0
|
||||
else 0.0
|
||||
),
|
||||
}
|
||||
return out
|
||||
|
|
@ -259,24 +263,30 @@ def _print_latency(title: str, lat: dict[str, dict]) -> None:
|
|||
print()
|
||||
print(title)
|
||||
print("-" * len(title))
|
||||
header = (f"{'arm':<25} {'n':>4} {'mean':>7} {'std':>7} "
|
||||
f"{'p50':>7} {'p90':>7} {'p95':>7} {'p99':>7} {'max':>7} {'CV':>5}")
|
||||
header = (
|
||||
f"{'arm':<25} {'n':>4} {'mean':>7} {'std':>7} "
|
||||
f"{'p50':>7} {'p90':>7} {'p95':>7} {'p99':>7} {'max':>7} {'CV':>5}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for arm in sorted(lat, key=lambda a: lat[a]["mean_s"]):
|
||||
s = lat[arm]
|
||||
print(f"{arm:<25} {s['n']:>4} "
|
||||
f"{s['mean_s']:>6.1f}s {s['std_s']:>6.1f}s "
|
||||
f"{s['p50_s']:>6.1f}s {s['p90_s']:>6.1f}s {s['p95_s']:>6.1f}s "
|
||||
f"{s['p99_s']:>6.1f}s {s['max_s']:>6.1f}s {s['cv']:>5.2f}")
|
||||
print(
|
||||
f"{arm:<25} {s['n']:>4} "
|
||||
f"{s['mean_s']:>6.1f}s {s['std_s']:>6.1f}s "
|
||||
f"{s['p50_s']:>6.1f}s {s['p90_s']:>6.1f}s {s['p95_s']:>6.1f}s "
|
||||
f"{s['p99_s']:>6.1f}s {s['max_s']:>6.1f}s {s['cv']:>5.2f}"
|
||||
)
|
||||
|
||||
|
||||
def _print_tokens(title: str, toks: dict[str, dict]) -> None:
|
||||
print()
|
||||
print(title)
|
||||
print("-" * len(title))
|
||||
header = (f"{'arm':<25} {'in mean':>9} {'in p50':>9} {'in p95':>9} {'in max':>9}"
|
||||
f" {'out mean':>9} {'out p95':>9}")
|
||||
header = (
|
||||
f"{'arm':<25} {'in mean':>9} {'in p50':>9} {'in p95':>9} {'in max':>9}"
|
||||
f" {'out mean':>9} {'out p95':>9}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for arm in sorted(toks):
|
||||
|
|
@ -285,25 +295,31 @@ def _print_tokens(title: str, toks: dict[str, dict]) -> None:
|
|||
eout = e.get("output")
|
||||
if not ein:
|
||||
continue
|
||||
print(f"{arm:<25} "
|
||||
f"{ein['mean']:>9,.0f} {ein['p50']:>9,.0f} {ein['p95']:>9,.0f} {ein['max']:>9,.0f} "
|
||||
f"{(eout or {}).get('mean', 0):>9,.0f} {(eout or {}).get('p95', 0):>9,.0f}")
|
||||
print(
|
||||
f"{arm:<25} "
|
||||
f"{ein['mean']:>9,.0f} {ein['p50']:>9,.0f} {ein['p95']:>9,.0f} {ein['max']:>9,.0f} "
|
||||
f"{(eout or {}).get('mean', 0):>9,.0f} {(eout or {}).get('p95', 0):>9,.0f}"
|
||||
)
|
||||
|
||||
|
||||
def _print_pdf_var(title: str, var: dict[str, dict]) -> None:
|
||||
print()
|
||||
print(title)
|
||||
print("-" * len(title))
|
||||
header = (f"{'arm':<25} {'n_pdfs':>7} {'mean':>7} {'std':>7} {'min':>7} "
|
||||
f"{'p25':>7} {'p50':>7} {'p75':>7} {'max':>7} {'#0%':>5} {'#100%':>6}")
|
||||
header = (
|
||||
f"{'arm':<25} {'n_pdfs':>7} {'mean':>7} {'std':>7} {'min':>7} "
|
||||
f"{'p25':>7} {'p50':>7} {'p75':>7} {'max':>7} {'#0%':>5} {'#100%':>6}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for arm in sorted(var, key=lambda a: -var[a]["mean"]):
|
||||
s = var[arm]
|
||||
print(f"{arm:<25} {s['n_pdfs']:>7} "
|
||||
f"{s['mean']*100:>6.1f}% {s['std']*100:>6.1f}% {s['min']*100:>6.1f}% "
|
||||
f"{s['p25']*100:>6.1f}% {s['p50']*100:>6.1f}% {s['p75']*100:>6.1f}% "
|
||||
f"{s['max']*100:>6.1f}% {s['n_pdfs_zero']:>5} {s['n_pdfs_perfect']:>6}")
|
||||
print(
|
||||
f"{arm:<25} {s['n_pdfs']:>7} "
|
||||
f"{s['mean'] * 100:>6.1f}% {s['std'] * 100:>6.1f}% {s['min'] * 100:>6.1f}% "
|
||||
f"{s['p25'] * 100:>6.1f}% {s['p50'] * 100:>6.1f}% {s['p75'] * 100:>6.1f}% "
|
||||
f"{s['max'] * 100:>6.1f}% {s['n_pdfs_zero']:>5} {s['n_pdfs_perfect']:>6}"
|
||||
)
|
||||
|
||||
|
||||
def _print_mcnemar(title: str, table: dict) -> None:
|
||||
|
|
@ -311,8 +327,10 @@ def _print_mcnemar(title: str, table: dict) -> None:
|
|||
print(title)
|
||||
print("-" * len(title))
|
||||
print(f"n_qids on which all arms have a graded row: {table['n_qids']}")
|
||||
header = (f"{'arm_i':<25} {'arm_j':<25} {'b':>4} {'c':>4} "
|
||||
f"{'both ok':>8} {'both wr':>8} {'p (2-sided)':>13} {'sig':>4}")
|
||||
header = (
|
||||
f"{'arm_i':<25} {'arm_j':<25} {'b':>4} {'c':>4} "
|
||||
f"{'both ok':>8} {'both wr':>8} {'p (2-sided)':>13} {'sig':>4}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for pair in sorted(table["pairs"], key=lambda p: p["p_value"]):
|
||||
|
|
@ -323,10 +341,12 @@ def _print_mcnemar(title: str, table: dict) -> None:
|
|||
sig = "**"
|
||||
elif pair["p_value"] < 0.05:
|
||||
sig = "*"
|
||||
print(f"{pair['arm_i']:<25} {pair['arm_j']:<25} "
|
||||
f"{pair['b_i_only']:>4} {pair['c_j_only']:>4} "
|
||||
f"{pair['both_correct']:>8} {pair['both_wrong']:>8} "
|
||||
f"{pair['p_value']:>13.4f} {sig:>4}")
|
||||
print(
|
||||
f"{pair['arm_i']:<25} {pair['arm_j']:<25} "
|
||||
f"{pair['b_i_only']:>4} {pair['c_j_only']:>4} "
|
||||
f"{pair['both_correct']:>8} {pair['both_wrong']:>8} "
|
||||
f"{pair['p_value']:>13.4f} {sig:>4}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -78,9 +78,11 @@ def _print_table(title: str, summary: dict[str, dict]) -> None:
|
|||
# stable order: highest accuracy first
|
||||
arms_sorted = sorted(summary.items(), key=lambda kv: -kv[1]["accuracy"])
|
||||
for arm, s in arms_sorted:
|
||||
print(f"{arm:<25} {s['n']:>4} {s['n_correct']:>7} "
|
||||
f"{s['accuracy']*100:>6.1f}% {s['f1_mean']*100:>6.1f}% "
|
||||
f"{s['n_failures']:>6} {s['failure_rate']*100:>6.1f}%")
|
||||
print(
|
||||
f"{arm:<25} {s['n']:>4} {s['n_correct']:>7} "
|
||||
f"{s['accuracy'] * 100:>6.1f}% {s['f1_mean'] * 100:>6.1f}% "
|
||||
f"{s['n_failures']:>6} {s['failure_rate'] * 100:>6.1f}%"
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
|
|
@ -103,9 +105,7 @@ def main() -> int:
|
|||
raw_rows = _read_jsonl(raw_path)
|
||||
retry_rows = _read_jsonl(retry_path)
|
||||
|
||||
retry_by_key: dict[tuple[str, str], dict] = {
|
||||
_row_key(r): r for r in retry_rows
|
||||
}
|
||||
retry_by_key: dict[tuple[str, str], dict] = {_row_key(r): r for r in retry_rows}
|
||||
|
||||
merged_rows: list[dict] = []
|
||||
n_replaced_recovered = 0
|
||||
|
|
|
|||
|
|
@ -44,10 +44,7 @@ def main() -> None:
|
|||
f"questions covering first 30 docs: total={len(qs_in_30)} "
|
||||
f"answerable={answerable} unanswerable={unanswerable}"
|
||||
)
|
||||
print(
|
||||
f"avg Qs/PDF: {len(qs_in_30) / 30:.1f} "
|
||||
f"answerable/PDF: {answerable / 30:.1f}"
|
||||
)
|
||||
print(f"avg Qs/PDF: {len(qs_in_30) / 30:.1f} answerable/PDF: {answerable / 30:.1f}")
|
||||
print(f"format mix in scope: {dict(fmts)}")
|
||||
print()
|
||||
print("25 new PDFs to ingest:")
|
||||
|
|
|
|||
|
|
@ -27,10 +27,7 @@ def main() -> None:
|
|||
grade = a.get("graded", {})
|
||||
text = (a.get("raw_text") or "").strip()
|
||||
tail = text[-200:] if text else ""
|
||||
print(
|
||||
f" [{arm_name}] grade={grade.get('grade')} "
|
||||
f"method={grade.get('method')}"
|
||||
)
|
||||
print(f" [{arm_name}] grade={grade.get('grade')} method={grade.get('method')}")
|
||||
print(f" -> {tail!r}")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -132,17 +132,19 @@ def _load_failed_rows(raw_path: Path) -> list[FailedRow]:
|
|||
row = json.loads(line)
|
||||
if not _is_failure_row(row):
|
||||
continue
|
||||
out.append(FailedRow(
|
||||
arm=str(row["arm"]),
|
||||
qid=str(row["qid"]),
|
||||
doc_id=str(row["doc_id"]),
|
||||
answer_format=str(row.get("answer_format") or ""),
|
||||
gold=str(row.get("gold") or ""),
|
||||
pages=int(row.get("pages") or 0),
|
||||
document_id=row.get("document_id"),
|
||||
original_error=row.get("error"),
|
||||
original_row=row,
|
||||
))
|
||||
out.append(
|
||||
FailedRow(
|
||||
arm=str(row["arm"]),
|
||||
qid=str(row["qid"]),
|
||||
doc_id=str(row["doc_id"]),
|
||||
answer_format=str(row.get("answer_format") or ""),
|
||||
gold=str(row.get("gold") or ""),
|
||||
pages=int(row.get("pages") or 0),
|
||||
document_id=row.get("document_id"),
|
||||
original_error=row.get("error"),
|
||||
original_row=row,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
|
|
@ -202,8 +204,12 @@ def _qid_index(qid: str) -> int:
|
|||
|
||||
|
||||
def _build_native_request(
|
||||
qid: str, question: str, answer_format: str, pdf_path: Path,
|
||||
*, max_output_tokens: int,
|
||||
qid: str,
|
||||
question: str,
|
||||
answer_format: str,
|
||||
pdf_path: Path,
|
||||
*,
|
||||
max_output_tokens: int,
|
||||
) -> ArmRequest:
|
||||
return ArmRequest(
|
||||
question_id=qid,
|
||||
|
|
@ -214,12 +220,14 @@ def _build_native_request(
|
|||
|
||||
|
||||
def _build_lc_request(
|
||||
qid: str, question: str, answer_format: str, doc_id: str, md_path: Path,
|
||||
qid: str,
|
||||
question: str,
|
||||
answer_format: str,
|
||||
doc_id: str,
|
||||
md_path: Path,
|
||||
) -> ArmRequest:
|
||||
if not md_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Missing parser extraction at {md_path}; cannot retry LC arm."
|
||||
)
|
||||
raise FileNotFoundError(f"Missing parser extraction at {md_path}; cannot retry LC arm.")
|
||||
markdown = md_path.read_text(encoding="utf-8")
|
||||
return ArmRequest(
|
||||
question_id=qid,
|
||||
|
|
@ -256,7 +264,9 @@ class RetryOutcome:
|
|||
|
||||
|
||||
async def _retry_one(
|
||||
arm_obj: Any, request: ArmRequest, *,
|
||||
arm_obj: Any,
|
||||
request: ArmRequest,
|
||||
*,
|
||||
arm_name: str,
|
||||
qid: str,
|
||||
max_attempts: int,
|
||||
|
|
@ -274,31 +284,44 @@ async def _retry_one(
|
|||
attempt_error = result.error
|
||||
if not attempt_error and not raw_text:
|
||||
attempt_error = "EmptyResponse: stream ended with no text"
|
||||
attempts.append(AttemptLog(
|
||||
attempt=attempt,
|
||||
started_iso=started_iso,
|
||||
latency_ms=latency_ms,
|
||||
error=attempt_error,
|
||||
raw_text_chars=len(raw_text),
|
||||
))
|
||||
attempts.append(
|
||||
AttemptLog(
|
||||
attempt=attempt,
|
||||
started_iso=started_iso,
|
||||
latency_ms=latency_ms,
|
||||
error=attempt_error,
|
||||
raw_text_chars=len(raw_text),
|
||||
)
|
||||
)
|
||||
final = result
|
||||
if not attempt_error and raw_text:
|
||||
return RetryOutcome(
|
||||
arm=arm_name, qid=qid, attempts=attempts,
|
||||
final_result=result, recovered=True,
|
||||
arm=arm_name,
|
||||
qid=qid,
|
||||
attempts=attempts,
|
||||
final_result=result,
|
||||
recovered=True,
|
||||
)
|
||||
if attempt < max_attempts:
|
||||
delay = min(max_delay, base_delay * (2 ** (attempt - 1)))
|
||||
delay = delay * (0.5 + random.random())
|
||||
logger.info(
|
||||
"[%s::%s] attempt %d/%d failed (%s); sleeping %.1fs",
|
||||
arm_name, qid, attempt, max_attempts, attempt_error, delay,
|
||||
arm_name,
|
||||
qid,
|
||||
attempt,
|
||||
max_attempts,
|
||||
attempt_error,
|
||||
delay,
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
assert final is not None
|
||||
return RetryOutcome(
|
||||
arm=arm_name, qid=qid, attempts=attempts,
|
||||
final_result=final, recovered=False,
|
||||
arm=arm_name,
|
||||
qid=qid,
|
||||
attempts=attempts,
|
||||
final_result=final,
|
||||
recovered=False,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -365,7 +388,8 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
by_arm_count[f.arm] = by_arm_count.get(f.arm, 0) + 1
|
||||
logger.info(
|
||||
"Loaded %d failed rows across %d arms: %s",
|
||||
len(failed), len(by_arm_count),
|
||||
len(failed),
|
||||
len(by_arm_count),
|
||||
", ".join(f"{a}={n}" for a, n in sorted(by_arm_count.items())),
|
||||
)
|
||||
|
||||
|
|
@ -383,7 +407,8 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
engine=PdfEngine(args.pdf_engine),
|
||||
)
|
||||
native_arm = NativePdfArm(
|
||||
provider=native_provider, max_output_tokens=args.max_output_tokens,
|
||||
provider=native_provider,
|
||||
max_output_tokens=args.max_output_tokens,
|
||||
)
|
||||
|
||||
lc_arms: dict[str, BareLlmArm] = {}
|
||||
|
|
@ -413,7 +438,8 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
if qrow is None:
|
||||
logger.error(
|
||||
"Could not find question text for %s (idx %d) — skipping",
|
||||
f.doc_id, q_idx,
|
||||
f.doc_id,
|
||||
q_idx,
|
||||
)
|
||||
continue
|
||||
question_text = str(qrow.get("question") or "").strip()
|
||||
|
|
@ -430,7 +456,10 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
logger.error("PDF missing on disk: %s — skipping", pdf_path)
|
||||
continue
|
||||
request = _build_native_request(
|
||||
f.qid, question_text, answer_format, pdf_path,
|
||||
f.qid,
|
||||
question_text,
|
||||
answer_format,
|
||||
pdf_path,
|
||||
max_output_tokens=args.max_output_tokens,
|
||||
)
|
||||
arm_obj = native_arm
|
||||
|
|
@ -440,11 +469,16 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
if not md_path_str or ext_blob.get("status") != "ok":
|
||||
logger.error(
|
||||
"Missing extraction for %s on %s — cannot retry; skipping",
|
||||
f.arm, f.doc_id,
|
||||
f.arm,
|
||||
f.doc_id,
|
||||
)
|
||||
continue
|
||||
request = _build_lc_request(
|
||||
f.qid, question_text, answer_format, f.doc_id, Path(md_path_str),
|
||||
f.qid,
|
||||
question_text,
|
||||
answer_format,
|
||||
f.doc_id,
|
||||
Path(md_path_str),
|
||||
)
|
||||
arm_obj = lc_arms[f.arm]
|
||||
else:
|
||||
|
|
@ -452,13 +486,17 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
continue
|
||||
|
||||
plan.append((f, request, arm_obj))
|
||||
coros.append(_retry_one(
|
||||
arm_obj, request,
|
||||
arm_name=f.arm, qid=f.qid,
|
||||
max_attempts=args.max_attempts,
|
||||
base_delay=args.base_delay,
|
||||
max_delay=args.max_delay,
|
||||
))
|
||||
coros.append(
|
||||
_retry_one(
|
||||
arm_obj,
|
||||
request,
|
||||
arm_name=f.arm,
|
||||
qid=f.qid,
|
||||
max_attempts=args.max_attempts,
|
||||
base_delay=args.base_delay,
|
||||
max_delay=args.max_delay,
|
||||
)
|
||||
)
|
||||
|
||||
if not coros:
|
||||
logger.warning("Nothing to retry after request building.")
|
||||
|
|
@ -467,13 +505,17 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
logger.info(
|
||||
"Retrying %d failed rows with up to %d attempts each "
|
||||
"(base_delay=%.1fs, max_delay=%.1fs, concurrency=%d).",
|
||||
len(coros), args.max_attempts, args.base_delay, args.max_delay,
|
||||
len(coros),
|
||||
args.max_attempts,
|
||||
args.base_delay,
|
||||
args.max_delay,
|
||||
args.concurrency,
|
||||
)
|
||||
|
||||
started = time.monotonic()
|
||||
outcomes: list[RetryOutcome] = await _gather_with_limit(
|
||||
coros, concurrency=args.concurrency,
|
||||
coros,
|
||||
concurrency=args.concurrency,
|
||||
)
|
||||
elapsed = time.monotonic() - started
|
||||
logger.info("Retry pass finished in %.1fs.", elapsed)
|
||||
|
|
@ -489,12 +531,8 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
for (f, _req, _arm_obj), outcome in zip(plan, outcomes, strict=True):
|
||||
per_arm_total[outcome.arm] = per_arm_total.get(outcome.arm, 0) + 1
|
||||
if outcome.recovered:
|
||||
per_arm_recovered[outcome.arm] = (
|
||||
per_arm_recovered.get(outcome.arm, 0) + 1
|
||||
)
|
||||
per_arm_attempts_dist.setdefault(outcome.arm, []).append(
|
||||
len(outcome.attempts)
|
||||
)
|
||||
per_arm_recovered[outcome.arm] = per_arm_recovered.get(outcome.arm, 0) + 1
|
||||
per_arm_attempts_dist.setdefault(outcome.arm, []).append(len(outcome.attempts))
|
||||
|
||||
g = grade(
|
||||
pred=extract_freeform_answer(outcome.final_result.raw_text or ""),
|
||||
|
|
@ -555,12 +593,11 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
arm: {
|
||||
"tried": per_arm_total.get(arm, 0),
|
||||
"recovered": per_arm_recovered.get(arm, 0),
|
||||
"still_failed": (
|
||||
per_arm_total.get(arm, 0) - per_arm_recovered.get(arm, 0)
|
||||
),
|
||||
"still_failed": (per_arm_total.get(arm, 0) - per_arm_recovered.get(arm, 0)),
|
||||
"recovery_rate": (
|
||||
per_arm_recovered.get(arm, 0) / per_arm_total[arm]
|
||||
if per_arm_total.get(arm) else 0.0
|
||||
if per_arm_total.get(arm)
|
||||
else 0.0
|
||||
),
|
||||
"attempts_distribution": sorted(per_arm_attempts_dist.get(arm, [])),
|
||||
}
|
||||
|
|
@ -593,8 +630,7 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
rec_total = sum(per_arm_recovered.values())
|
||||
rate_total = (rec_total / total * 100) if total else 0.0
|
||||
print("-" * len(header))
|
||||
print(f"{'TOTAL':<25} {total:>6} {rec_total:>10} {total - rec_total:>11} "
|
||||
f"{rate_total:>6.1f}%")
|
||||
print(f"{'TOTAL':<25} {total:>6} {rec_total:>10} {total - rec_total:>11} {rate_total:>6.1f}%")
|
||||
print()
|
||||
print(f"Wrote {out_path.relative_to(REPO)}")
|
||||
print(f"Wrote {summary_path.relative_to(REPO)}")
|
||||
|
|
@ -604,27 +640,37 @@ async def _run(args: argparse.Namespace) -> int:
|
|||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--run-id", default="2026-05-14T00-53-19Z",
|
||||
"--run-id",
|
||||
default="2026-05-14T00-53-19Z",
|
||||
help="Run timestamp under data/multimodal_doc/runs/. Default is the "
|
||||
"n=171 production run we wrote up in the blog.",
|
||||
"n=171 production run we wrote up in the blog.",
|
||||
)
|
||||
parser.add_argument("--max-attempts", type=int, default=5)
|
||||
parser.add_argument("--base-delay", type=float, default=1.0,
|
||||
help="Base seconds for exponential backoff (default 1s).")
|
||||
parser.add_argument("--max-delay", type=float, default=30.0,
|
||||
help="Cap on per-retry sleep (default 30s).")
|
||||
parser.add_argument("--concurrency", type=int, default=2,
|
||||
help="Parallel retries in flight (default 2 — keep low "
|
||||
"to avoid the same transport stress that caused "
|
||||
"the original failures).")
|
||||
parser.add_argument(
|
||||
"--base-delay",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Base seconds for exponential backoff (default 1s).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-delay", type=float, default=30.0, help="Cap on per-retry sleep (default 30s)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Parallel retries in flight (default 2 — keep low "
|
||||
"to avoid the same transport stress that caused "
|
||||
"the original failures).",
|
||||
)
|
||||
parser.add_argument("--llm-model", default="anthropic/claude-sonnet-4.5")
|
||||
parser.add_argument("--pdf-engine", default="native",
|
||||
choices=[e.value for e in PdfEngine])
|
||||
parser.add_argument("--pdf-engine", default="native", choices=[e.value for e in PdfEngine])
|
||||
parser.add_argument("--max-output-tokens", type=int, default=512)
|
||||
parser.add_argument(
|
||||
"--include-surfsense", action="store_true",
|
||||
"--include-surfsense",
|
||||
action="store_true",
|
||||
help="Also retry surfsense_agentic failures (requires backend + celery up). "
|
||||
"Default is to skip them since the n=171 run had 0 SurfSense failures.",
|
||||
"Default is to skip them since the n=171 run had 0 SurfSense failures.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
raise SystemExit(asyncio.run(_run(args)))
|
||||
|
|
|
|||
|
|
@ -23,12 +23,12 @@ def main() -> None:
|
|||
d = metrics[arm]
|
||||
print(
|
||||
f"{arm:14s}: "
|
||||
f"acc={d['accuracy']*100:5.1f}% (Wilson 95% CI "
|
||||
f"{d['ci_low']*100:.1f}-{d['ci_high']*100:.1f}) | "
|
||||
f"correct={d['correct_rate']*100:5.1f}% "
|
||||
f"missing={d['missing_rate']*100:5.1f}% "
|
||||
f"incorrect={d['incorrect_rate']*100:5.1f}% | "
|
||||
f"truth={d['truthfulness_score']*100:+5.1f}%"
|
||||
f"acc={d['accuracy'] * 100:5.1f}% (Wilson 95% CI "
|
||||
f"{d['ci_low'] * 100:.1f}-{d['ci_high'] * 100:.1f}) | "
|
||||
f"correct={d['correct_rate'] * 100:5.1f}% "
|
||||
f"missing={d['missing_rate'] * 100:5.1f}% "
|
||||
f"incorrect={d['incorrect_rate'] * 100:5.1f}% | "
|
||||
f"truth={d['truthfulness_score'] * 100:+5.1f}%"
|
||||
)
|
||||
|
||||
print()
|
||||
|
|
@ -48,7 +48,7 @@ def main() -> None:
|
|||
pieces = [f"{qt:20s} (n={n:3d}):"]
|
||||
for arm in ("bare_llm", "long_context", "surfsense"):
|
||||
if arm in row:
|
||||
pieces.append(f"{arm}={row[arm]['truthfulness_score']*100:+7.1f}%")
|
||||
pieces.append(f"{arm}={row[arm]['truthfulness_score'] * 100:+7.1f}%")
|
||||
print(" ".join(pieces))
|
||||
|
||||
print()
|
||||
|
|
@ -58,7 +58,7 @@ def main() -> None:
|
|||
pieces = [f"{dom:10s} (n={n:3d}):"]
|
||||
for arm in ("bare_llm", "long_context", "surfsense"):
|
||||
if arm in row:
|
||||
pieces.append(f"{arm}={row[arm]['truthfulness_score']*100:+7.1f}%")
|
||||
pieces.append(f"{arm}={row[arm]['truthfulness_score'] * 100:+7.1f}%")
|
||||
print(" ".join(pieces))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -23,7 +23,9 @@ ARTIFACT = RUN_DIR / "run_artifact.json"
|
|||
|
||||
|
||||
def main() -> None:
|
||||
rows = [json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()]
|
||||
rows = [
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
|
||||
]
|
||||
print(f"raw rows: {len(rows)}")
|
||||
|
||||
by_qid: dict[str, list[dict]] = defaultdict(list)
|
||||
|
|
@ -31,11 +33,19 @@ def main() -> None:
|
|||
by_qid[row["qid"]].append(row)
|
||||
print(f"unique questions: {len(by_qid)}")
|
||||
|
||||
arm_metrics: dict[str, dict] = defaultdict(lambda: {
|
||||
"n": 0, "n_correct": 0, "n_failed": 0, "n_empty": 0,
|
||||
"costs": [], "in_tokens": [], "out_tokens": [], "latency_ms": [],
|
||||
"by_format": defaultdict(lambda: {"n": 0, "correct": 0}),
|
||||
})
|
||||
arm_metrics: dict[str, dict] = defaultdict(
|
||||
lambda: {
|
||||
"n": 0,
|
||||
"n_correct": 0,
|
||||
"n_failed": 0,
|
||||
"n_empty": 0,
|
||||
"costs": [],
|
||||
"in_tokens": [],
|
||||
"out_tokens": [],
|
||||
"latency_ms": [],
|
||||
"by_format": defaultdict(lambda: {"n": 0, "correct": 0}),
|
||||
}
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
arm = row["arm"]
|
||||
|
|
@ -70,7 +80,9 @@ def main() -> None:
|
|||
|
||||
print()
|
||||
print("=" * 100)
|
||||
print(f"{'arm':<25} {'n':>4} {'acc%':>6} {'F1%':>6} {'fail':>5} {'$ mean':>10} {'$ median':>10} {'in tok mean':>12} {'out tok mean':>12} {'p50 ms':>8}")
|
||||
print(
|
||||
f"{'arm':<25} {'n':>4} {'acc%':>6} {'F1%':>6} {'fail':>5} {'$ mean':>10} {'$ median':>10} {'in tok mean':>12} {'out tok mean':>12} {'p50 ms':>8}"
|
||||
)
|
||||
print("=" * 100)
|
||||
art = json.loads(ARTIFACT.read_text(encoding="utf-8"))
|
||||
per_arm_art = art["metrics"]["per_arm"]
|
||||
|
|
@ -110,7 +122,7 @@ def main() -> None:
|
|||
print("Aggregated cost (from run_artifact.json):")
|
||||
for arm, row in per_arm_art.items():
|
||||
print(
|
||||
f" {arm:<25} acc={row['accuracy']*100:5.1f}% "
|
||||
f" {arm:<25} acc={row['accuracy'] * 100:5.1f}% "
|
||||
f" $/Q LLM={row['llm_cost_per_q']:.4f} "
|
||||
f" preprocess total=${row['preprocess_cost_total']:.2f} "
|
||||
f" $/Q total={row['total_cost_per_q']:.4f}"
|
||||
|
|
|
|||
|
|
@ -40,8 +40,7 @@ CONTEXT_HINTS = (
|
|||
|
||||
def main() -> None:
|
||||
rows = [
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip()
|
||||
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
|
||||
]
|
||||
|
||||
extraction_size: dict[tuple[str, str], int] = {}
|
||||
|
|
@ -73,12 +72,12 @@ def main() -> None:
|
|||
print("=" * 80)
|
||||
print("(b) Extraction size for OK vs FAILED rows per arm")
|
||||
print("=" * 80)
|
||||
arm_buckets: dict[str, dict[str, list[int]]] = defaultdict(
|
||||
lambda: {"ok": [], "fail": []}
|
||||
)
|
||||
arm_buckets: dict[str, dict[str, list[int]]] = defaultdict(lambda: {"ok": [], "fail": []})
|
||||
parser_arms = (
|
||||
"azure_basic_lc", "azure_premium_lc",
|
||||
"llamacloud_basic_lc", "llamacloud_premium_lc",
|
||||
"azure_basic_lc",
|
||||
"azure_premium_lc",
|
||||
"llamacloud_basic_lc",
|
||||
"llamacloud_premium_lc",
|
||||
)
|
||||
for row in rows:
|
||||
arm = row["arm"]
|
||||
|
|
@ -133,10 +132,13 @@ def main() -> None:
|
|||
" 3M_2018_10K x llamacloud_premium = 908,733 chars (~227k tokens) "
|
||||
"-- this is above Sonnet 4.5's 200k window."
|
||||
)
|
||||
print(" If transport hypothesis is correct, this should still fail with a "
|
||||
"real overflow error.")
|
||||
print(" If transport hypothesis is correct AND the model truncates silently, "
|
||||
"it might 'succeed' but be wrong.")
|
||||
print(
|
||||
" If transport hypothesis is correct, this should still fail with a real overflow error."
|
||||
)
|
||||
print(
|
||||
" If transport hypothesis is correct AND the model truncates silently, "
|
||||
"it might 'succeed' but be wrong."
|
||||
)
|
||||
print()
|
||||
for row in rows:
|
||||
if row["doc_id"] != "3M_2018_10K.pdf":
|
||||
|
|
@ -145,10 +147,7 @@ def main() -> None:
|
|||
continue
|
||||
err = row.get("error") or "(none)"
|
||||
graded = row.get("graded") or {}
|
||||
print(
|
||||
f" {row['qid']:<40} correct={graded.get('correct')!s:<5} "
|
||||
f"err={err[:100]}"
|
||||
)
|
||||
print(f" {row['qid']:<40} correct={graded.get('correct')!s:<5} err={err[:100]}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -72,9 +72,7 @@ class SurfSenseArm(Arm):
|
|||
try:
|
||||
await self._client.delete_thread(thread_id)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug(
|
||||
"Failed to delete thread %s: %s", thread_id, exc
|
||||
)
|
||||
logger.debug("Failed to delete thread %s: %s", thread_id, exc)
|
||||
|
||||
letter = extract_answer_letter(answer.text)
|
||||
return ArmResult(
|
||||
|
|
|
|||
|
|
@ -83,6 +83,7 @@ async def acquire_token(config: Config, *, http: httpx.AsyncClient | None = None
|
|||
)
|
||||
|
||||
if config.has_local_mode():
|
||||
|
||||
async def _login(client: httpx.AsyncClient) -> TokenBundle:
|
||||
response = await client.post(
|
||||
f"{config.surfsense_api_base}/auth/desktop/login",
|
||||
|
|
@ -94,15 +95,12 @@ async def acquire_token(config: Config, *, http: httpx.AsyncClient | None = None
|
|||
)
|
||||
if response.status_code != 200:
|
||||
raise CredentialError(
|
||||
f"LOCAL login failed (HTTP {response.status_code}): "
|
||||
f"{_safe_text(response)}"
|
||||
f"LOCAL login failed (HTTP {response.status_code}): {_safe_text(response)}"
|
||||
)
|
||||
payload = response.json()
|
||||
access = payload.get("access_token")
|
||||
if not access:
|
||||
raise CredentialError(
|
||||
f"LOCAL login response missing access_token: {payload!r}"
|
||||
)
|
||||
raise CredentialError(f"LOCAL login response missing access_token: {payload!r}")
|
||||
return TokenBundle(
|
||||
access_token=access,
|
||||
refresh_token=payload.get("refresh_token") or None,
|
||||
|
|
|
|||
|
|
@ -204,8 +204,7 @@ async def _cmd_setup(args: argparse.Namespace) -> int:
|
|||
|
||||
if scenario not in SCENARIOS:
|
||||
console.print(
|
||||
f"[red]Unknown scenario {scenario!r}. Pick one of: "
|
||||
f"{', '.join(SCENARIOS)}[/red]"
|
||||
f"[red]Unknown scenario {scenario!r}. Pick one of: {', '.join(SCENARIOS)}[/red]"
|
||||
)
|
||||
return 2
|
||||
|
||||
|
|
@ -292,9 +291,7 @@ async def _cmd_setup(args: argparse.Namespace) -> int:
|
|||
if not skip_vision_setup and (vision_required or vision_llm_slug is not None):
|
||||
try:
|
||||
vision_candidates = await ss_client.list_global_vision_models()
|
||||
resolved = resolve_vision_llm(
|
||||
vision_candidates, explicit_slug=vision_llm_slug
|
||||
)
|
||||
resolved = resolve_vision_llm(vision_candidates, explicit_slug=vision_llm_slug)
|
||||
except VisionConfigError as exc:
|
||||
console.print(f"[red]{exc}[/red]")
|
||||
return 2
|
||||
|
|
@ -524,10 +521,7 @@ async def _cmd_run(args: argparse.Namespace) -> int:
|
|||
)
|
||||
artifact = await benchmark.run(ctx, **extra_kwargs)
|
||||
|
||||
console.print(
|
||||
f"[green]run OK[/green] {args.suite}/{args.benchmark} → "
|
||||
f"{artifact.raw_path}"
|
||||
)
|
||||
console.print(f"[green]run OK[/green] {args.suite}/{args.benchmark} → {artifact.raw_path}")
|
||||
return 0
|
||||
|
||||
|
||||
|
|
@ -697,15 +691,21 @@ def _build_parser() -> argparse.ArgumentParser:
|
|||
)
|
||||
p_setup.set_defaults(_func=_cmd_setup, _async=True)
|
||||
|
||||
p_teardown = sub.add_parser("teardown", help="Soft-delete the suite SearchSpace + clear state slot.")
|
||||
p_teardown = sub.add_parser(
|
||||
"teardown", help="Soft-delete the suite SearchSpace + clear state slot."
|
||||
)
|
||||
p_teardown.add_argument("--suite", required=True)
|
||||
p_teardown.set_defaults(_func=_cmd_teardown, _async=True)
|
||||
|
||||
p_models = sub.add_parser("models", help="LLM-config discovery helpers.")
|
||||
models_sub = p_models.add_subparsers(dest="subcommand", required=True)
|
||||
p_models_list = models_sub.add_parser("list", help="List global LLM configs.")
|
||||
p_models_list.add_argument("--provider", default=None, help="Filter by provider, e.g. openrouter")
|
||||
p_models_list.add_argument("--grep", default=None, help="Substring filter on name / model_name.")
|
||||
p_models_list.add_argument(
|
||||
"--provider", default=None, help="Filter by provider, e.g. openrouter"
|
||||
)
|
||||
p_models_list.add_argument(
|
||||
"--grep", default=None, help="Substring filter on name / model_name."
|
||||
)
|
||||
p_models_list.set_defaults(_func=_cmd_models_list, _async=True)
|
||||
|
||||
p_suites = sub.add_parser("suites", help="List registered suites.")
|
||||
|
|
@ -729,7 +729,9 @@ def _build_parser() -> argparse.ArgumentParser:
|
|||
suite_parser = ingest_sub.add_parser(suite, help=f"Ingest a {suite} benchmark.")
|
||||
suite_bench = suite_parser.add_subparsers(dest="benchmark", required=True)
|
||||
for benchmark in registry.list_benchmarks(suite):
|
||||
bp = suite_bench.add_parser(benchmark.name, help=getattr(benchmark, "description", benchmark.name))
|
||||
bp = suite_bench.add_parser(
|
||||
benchmark.name, help=getattr(benchmark, "description", benchmark.name)
|
||||
)
|
||||
if hasattr(benchmark, "add_run_args"):
|
||||
benchmark.add_run_args(bp)
|
||||
bp.set_defaults(_func=_cmd_ingest, _async=True)
|
||||
|
|
@ -740,7 +742,9 @@ def _build_parser() -> argparse.ArgumentParser:
|
|||
suite_parser = run_sub.add_parser(suite, help=f"Run a {suite} benchmark.")
|
||||
suite_bench = suite_parser.add_subparsers(dest="benchmark", required=True)
|
||||
for benchmark in registry.list_benchmarks(suite):
|
||||
bp = suite_bench.add_parser(benchmark.name, help=getattr(benchmark, "description", benchmark.name))
|
||||
bp = suite_bench.add_parser(
|
||||
benchmark.name, help=getattr(benchmark, "description", benchmark.name)
|
||||
)
|
||||
if hasattr(benchmark, "add_run_args"):
|
||||
benchmark.add_run_args(bp)
|
||||
bp.set_defaults(_func=_cmd_run, _async=True)
|
||||
|
|
|
|||
|
|
@ -84,8 +84,7 @@ class DocumentProcessingFailed(RuntimeError):
|
|||
|
||||
def __init__(self, statuses: Sequence[DocumentStatus]) -> None:
|
||||
details = ", ".join(
|
||||
f"id={s.document_id} ({s.title!r}): {s.reason or 'unknown'}"
|
||||
for s in statuses
|
||||
f"id={s.document_id} ({s.title!r}): {s.reason or 'unknown'}" for s in statuses
|
||||
)
|
||||
super().__init__(f"Document(s) failed to process: {details}")
|
||||
self.statuses = list(statuses)
|
||||
|
|
@ -240,9 +239,7 @@ class DocumentsClient:
|
|||
# chunks (chunk_id -> document_id map)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def list_chunks(
|
||||
self, document_id: int, *, page_size: int = 100
|
||||
) -> list[ChunkRow]:
|
||||
async def list_chunks(self, document_id: int, *, page_size: int = 100) -> list[ChunkRow]:
|
||||
"""Walk ``GET /documents/{id}/chunks`` until ``has_more=False``.
|
||||
|
||||
Used by ingestion to materialise the ``chunk_id -> document_id``
|
||||
|
|
|
|||
|
|
@ -145,7 +145,7 @@ class NewChatClient:
|
|||
if attempt > max_busy_retries:
|
||||
raise
|
||||
# Cap wait at 30s; backend retry hint is exponential anyway.
|
||||
wait = min(30.0, 0.5 * (2 ** attempt))
|
||||
wait = min(30.0, 0.5 * (2**attempt))
|
||||
logger.info(
|
||||
"thread_id=%s busy (%s); retry %d/%d after %.1fs",
|
||||
thread_id,
|
||||
|
|
|
|||
|
|
@ -177,16 +177,12 @@ class SearchSpaceClient:
|
|||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
if not isinstance(payload, list):
|
||||
raise RuntimeError(
|
||||
f"Unexpected /model-connections/global payload: {payload!r}"
|
||||
)
|
||||
raise RuntimeError(f"Unexpected /model-connections/global payload: {payload!r}")
|
||||
entries: list[VisionModelEntry] = []
|
||||
for connection in payload:
|
||||
provider = str(connection.get("provider", ""))
|
||||
for model in connection.get("models") or []:
|
||||
if not model.get("enabled", True) or not model.get("supports_image_input"):
|
||||
continue
|
||||
entries.append(
|
||||
VisionModelEntry.from_payload({**model, "provider": provider})
|
||||
)
|
||||
entries.append(VisionModelEntry.from_payload({**model, "provider": provider}))
|
||||
return entries
|
||||
|
|
|
|||
|
|
@ -104,7 +104,9 @@ def load_config() -> Config:
|
|||
data_dir = Path(os.environ.get("EVAL_DATA_DIR") or (project_root / "data")).resolve()
|
||||
reports_dir = Path(os.environ.get("EVAL_REPORTS_DIR") or (project_root / "reports")).resolve()
|
||||
return Config(
|
||||
surfsense_api_base=os.environ.get("SURFSENSE_API_BASE", "http://localhost:8000").rstrip("/"),
|
||||
surfsense_api_base=os.environ.get("SURFSENSE_API_BASE", "http://localhost:8000").rstrip(
|
||||
"/"
|
||||
),
|
||||
openrouter_api_key=os.environ.get("OPENROUTER_API_KEY") or None,
|
||||
openrouter_base_url=os.environ.get(
|
||||
"OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"
|
||||
|
|
@ -203,9 +205,7 @@ class SuiteState:
|
|||
else None
|
||||
),
|
||||
native_arm_model=(
|
||||
str(payload["native_arm_model"])
|
||||
if payload.get("native_arm_model")
|
||||
else None
|
||||
str(payload["native_arm_model"]) if payload.get("native_arm_model") else None
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -95,10 +95,7 @@ class IngestSettings:
|
|||
def render_label(self) -> str:
|
||||
"""Human-readable single-line label for reports / log lines."""
|
||||
|
||||
return (
|
||||
f"vision={'on' if self.use_vision_llm else 'off'}, "
|
||||
f"mode={self.processing_mode}"
|
||||
)
|
||||
return f"vision={'on' if self.use_vision_llm else 'off'}, mode={self.processing_mode}"
|
||||
|
||||
|
||||
def _coerce_bool(value: Any, default: bool) -> bool:
|
||||
|
|
@ -122,9 +119,7 @@ def _coerce_mode(value: Any, default: str) -> str:
|
|||
return default
|
||||
val = str(value).strip().lower()
|
||||
if val not in PROCESSING_MODES:
|
||||
raise ValueError(
|
||||
f"Invalid processing_mode {val!r}; must be one of {PROCESSING_MODES}"
|
||||
)
|
||||
raise ValueError(f"Invalid processing_mode {val!r}; must be one of {PROCESSING_MODES}")
|
||||
return val
|
||||
|
||||
|
||||
|
|
@ -274,10 +269,7 @@ def format_ingest_settings_md(settings: Any) -> str:
|
|||
return "- SurfSense ingest settings: (not recorded — re-ingest to capture)"
|
||||
vision = "on" if settings.get("use_vision_llm") else "off"
|
||||
mode = settings.get("processing_mode") or "basic"
|
||||
return (
|
||||
f"- SurfSense ingest settings: vision_llm=`{vision}`, "
|
||||
f"processing_mode=`{mode}`"
|
||||
)
|
||||
return f"- SurfSense ingest settings: vision_llm=`{vision}`, processing_mode=`{mode}`"
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
|
|
|||
|
|
@ -67,17 +67,13 @@ def mcnemar_test(
|
|||
"""
|
||||
|
||||
if len(arm_a_correct) != len(arm_b_correct):
|
||||
raise ValueError(
|
||||
f"Length mismatch: arm_a={len(arm_a_correct)}, arm_b={len(arm_b_correct)}"
|
||||
)
|
||||
raise ValueError(f"Length mismatch: arm_a={len(arm_a_correct)}, arm_b={len(arm_b_correct)}")
|
||||
n = len(arm_a_correct)
|
||||
b = sum(1 for a, c in zip(arm_a_correct, arm_b_correct, strict=False) if a and not c)
|
||||
c = sum(1 for a, cc in zip(arm_a_correct, arm_b_correct, strict=False) if (not a) and cc)
|
||||
discordant = b + c
|
||||
if discordant == 0:
|
||||
return McnemarResult(
|
||||
n_total=n, b=b, c=c, statistic=0.0, p_value=1.0, method="degenerate"
|
||||
)
|
||||
return McnemarResult(n_total=n, b=b, c=c, statistic=0.0, p_value=1.0, method="degenerate")
|
||||
|
||||
if discordant < use_exact_below:
|
||||
# Exact binomial: under H0 each discordant pair is a Bernoulli(0.5).
|
||||
|
|
@ -92,13 +88,11 @@ def mcnemar_test(
|
|||
# Chi-square with continuity correction (McNemar-Edwards).
|
||||
chi = ((abs(b - c) - 1) ** 2) / discordant
|
||||
p_value = _chi2_sf(chi, df=1)
|
||||
return McnemarResult(
|
||||
n_total=n, b=b, c=c, statistic=chi, p_value=p_value, method="chi2_cc"
|
||||
)
|
||||
return McnemarResult(n_total=n, b=b, c=c, statistic=chi, p_value=p_value, method="chi2_cc")
|
||||
|
||||
|
||||
def _binom_pmf(n: int, k: int) -> float:
|
||||
return math.comb(n, k) * (0.5 ** n)
|
||||
return math.comb(n, k) * (0.5**n)
|
||||
|
||||
|
||||
def _chi2_sf(x: float, *, df: int) -> float:
|
||||
|
|
|
|||
|
|
@ -46,9 +46,7 @@ _Z_FOR_LEVEL: dict[float, float] = {
|
|||
}
|
||||
|
||||
|
||||
def wilson_ci(
|
||||
n_correct: int, n_total: int, *, level: float = 0.95
|
||||
) -> tuple[float, float]:
|
||||
def wilson_ci(n_correct: int, n_total: int, *, level: float = 0.95) -> tuple[float, float]:
|
||||
"""Two-sided Wilson score confidence interval for a proportion.
|
||||
|
||||
Returns ``(low, high)``. ``n_total == 0`` returns ``(0.0, 1.0)`` —
|
||||
|
|
@ -70,9 +68,7 @@ def wilson_ci(
|
|||
return low, high
|
||||
|
||||
|
||||
def accuracy_with_wilson_ci(
|
||||
n_correct: int, n_total: int, *, level: float = 0.95
|
||||
) -> AccuracyResult:
|
||||
def accuracy_with_wilson_ci(n_correct: int, n_total: int, *, level: float = 0.95) -> AccuracyResult:
|
||||
if n_total < 0:
|
||||
raise ValueError(f"n_total must be >= 0, got {n_total}")
|
||||
if n_correct < 0 or n_correct > n_total:
|
||||
|
|
@ -109,10 +105,7 @@ def per_task_accuracy(
|
|||
bucket[1] += 1
|
||||
if row.get(correct_key):
|
||||
bucket[0] += 1
|
||||
return {
|
||||
task: accuracy_with_wilson_ci(c[0], c[1], level=level)
|
||||
for task, c in counts.items()
|
||||
}
|
||||
return {task: accuracy_with_wilson_ci(c[0], c[1], level=level) for task, c in counts.items()}
|
||||
|
||||
|
||||
def macro_accuracy(per_task: Mapping[str, AccuracyResult]) -> float:
|
||||
|
|
|
|||
|
|
@ -61,7 +61,7 @@ def _dcg_at_k(grades: Sequence[float], k: int) -> float:
|
|||
s = 0.0
|
||||
for i, grade in enumerate(grades[:k], start=1):
|
||||
# Standard log-base-2 discount; gain = 2^grade - 1 for graded relevance.
|
||||
s += (2.0 ** grade - 1.0) / math.log2(i + 1)
|
||||
s += (2.0**grade - 1.0) / math.log2(i + 1)
|
||||
return s
|
||||
|
||||
|
||||
|
|
@ -106,7 +106,9 @@ def score_run(
|
|||
|
||||
qids = set(per_query_qrels.keys()) & set(per_query_retrieved.keys())
|
||||
if not qids:
|
||||
return RetrievalScores(recall_at_k={k: 0.0 for k in ks}, mrr=0.0, ndcg_at_10=0.0, n_queries=0)
|
||||
return RetrievalScores(
|
||||
recall_at_k={k: 0.0 for k in ks}, mrr=0.0, ndcg_at_10=0.0, n_queries=0
|
||||
)
|
||||
|
||||
recall_totals = {k: 0.0 for k in ks}
|
||||
mrr_total = 0.0
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ from typing import Any
|
|||
# the pattern source, so we splice the literal character in via an
|
||||
# f-string. This keeps our pattern functionally identical to the TS
|
||||
# reference and lets ``"\u200B" in CITATION_REGEX.pattern`` succeed.
|
||||
_ZWSP = "\u200B"
|
||||
_ZWSP = "\u200b"
|
||||
CITATION_REGEX = re.compile(
|
||||
rf"[\[【]{_ZWSP}?citation:\s*("
|
||||
rf"https?://[^\]】{_ZWSP}]+|urlcite\d+|(?:doc-)?-?\d+(?:\s*,\s*(?:doc-)?-?\d+)*"
|
||||
|
|
|
|||
|
|
@ -56,7 +56,7 @@ def extract_freeform_answer(text: str) -> str:
|
|||
marker_matches = list(_ANSWER_MARKER.finditer(text))
|
||||
if marker_matches:
|
||||
last = marker_matches[-1]
|
||||
tail = text[last.end():]
|
||||
tail = text[last.end() :]
|
||||
nl = tail.find("\n")
|
||||
if nl >= 0:
|
||||
tail = tail[:nl]
|
||||
|
|
@ -77,7 +77,7 @@ def extract_freeform_answer(text: str) -> str:
|
|||
# 2. Strip wrapping quotes / parens / trailing punctuation that
|
||||
# confuse the grader without changing meaning.
|
||||
candidate = candidate.strip().strip("`").strip()
|
||||
if candidate.startswith(("\"", "'")) and candidate.endswith(("\"", "'")):
|
||||
if candidate.startswith(('"', "'")) and candidate.endswith(('"', "'")):
|
||||
candidate = candidate[1:-1].strip()
|
||||
return candidate
|
||||
|
||||
|
|
|
|||
|
|
@ -64,8 +64,7 @@ async def parse_with_azure_di(
|
|||
api_key = api_key or os.environ.get("AZURE_DI_KEY")
|
||||
if not endpoint or not api_key:
|
||||
raise ValueError(
|
||||
"AZURE_DI_ENDPOINT and AZURE_DI_KEY must be set "
|
||||
"(see surfsense_evals/.env)."
|
||||
"AZURE_DI_ENDPOINT and AZURE_DI_KEY must be set (see surfsense_evals/.env)."
|
||||
)
|
||||
|
||||
model_id = _AZURE_MODEL_BY_MODE.get(processing_mode, "prebuilt-read")
|
||||
|
|
@ -86,7 +85,10 @@ async def parse_with_azure_di(
|
|||
file_size_mb = await asyncio.to_thread(os.path.getsize, file_path) / (1024 * 1024)
|
||||
logger.info(
|
||||
"Azure DI parsing %s (mode=%s, model=%s, size=%.1fMB)",
|
||||
file_path, processing_mode, model_id, file_size_mb,
|
||||
file_path,
|
||||
processing_mode,
|
||||
model_id,
|
||||
file_size_mb,
|
||||
)
|
||||
|
||||
last_exc: Exception | None = None
|
||||
|
|
@ -106,12 +108,12 @@ async def parse_with_azure_di(
|
|||
result = await poller.result()
|
||||
content = (result.content or "").strip()
|
||||
if not content:
|
||||
raise AzureDIError(
|
||||
f"Azure DI returned empty content for {file_path}"
|
||||
)
|
||||
raise AzureDIError(f"Azure DI returned empty content for {file_path}")
|
||||
logger.info(
|
||||
"Azure DI OK: %s (%s) -> %d chars",
|
||||
file_path, model_id, len(content),
|
||||
file_path,
|
||||
model_id,
|
||||
len(content),
|
||||
)
|
||||
return content
|
||||
|
||||
|
|
@ -120,9 +122,7 @@ async def parse_with_azure_di(
|
|||
except HttpResponseError as exc:
|
||||
# 4xx that's not auth: don't retry, the request itself is broken.
|
||||
if exc.status_code and 400 <= exc.status_code < 500:
|
||||
raise AzureDIError(
|
||||
f"Azure DI {exc.status_code} on {file_path}: {exc}"
|
||||
) from exc
|
||||
raise AzureDIError(f"Azure DI {exc.status_code} on {file_path}: {exc}") from exc
|
||||
last_exc = exc
|
||||
except (ServiceRequestError, ServiceResponseError) as exc:
|
||||
last_exc = exc
|
||||
|
|
@ -133,7 +133,10 @@ async def parse_with_azure_di(
|
|||
sleep_for = delay + jitter
|
||||
logger.warning(
|
||||
"Azure DI attempt %d/%d failed (%s); retrying in %.1fs",
|
||||
attempt, _MAX_RETRIES, type(last_exc).__name__, sleep_for,
|
||||
attempt,
|
||||
_MAX_RETRIES,
|
||||
type(last_exc).__name__,
|
||||
sleep_for,
|
||||
)
|
||||
await asyncio.sleep(sleep_for)
|
||||
|
||||
|
|
|
|||
|
|
@ -61,8 +61,7 @@ def _extract_markdown(result) -> str:
|
|||
if result and hasattr(result[0], "text"):
|
||||
return result[0].text
|
||||
return "\n\n".join(
|
||||
doc.page_content if hasattr(doc, "page_content") else str(doc)
|
||||
for doc in result
|
||||
doc.page_content if hasattr(doc, "page_content") else str(doc) for doc in result
|
||||
)
|
||||
|
||||
return str(result)
|
||||
|
|
@ -86,9 +85,7 @@ async def parse_with_llamacloud(
|
|||
|
||||
api_key = api_key or os.environ.get("LLAMA_CLOUD_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"LLAMA_CLOUD_API_KEY must be set (see surfsense_evals/.env)."
|
||||
)
|
||||
raise ValueError("LLAMA_CLOUD_API_KEY must be set (see surfsense_evals/.env).")
|
||||
|
||||
parse_mode = _LLAMA_PARSE_MODE_MAP.get(processing_mode, "parse_page_with_llm")
|
||||
|
||||
|
|
@ -106,13 +103,19 @@ async def parse_with_llamacloud(
|
|||
upload_timeout = max(120.0, 30.0 * file_size_mb)
|
||||
|
||||
logger.info(
|
||||
"LlamaCloud parsing %s (mode=%s, parse_mode=%s, %.1fMB, "
|
||||
"job_timeout=%.0fs)",
|
||||
file_path, processing_mode, parse_mode, file_size_mb, job_timeout,
|
||||
"LlamaCloud parsing %s (mode=%s, parse_mode=%s, %.1fMB, job_timeout=%.0fs)",
|
||||
file_path,
|
||||
processing_mode,
|
||||
parse_mode,
|
||||
file_size_mb,
|
||||
job_timeout,
|
||||
)
|
||||
|
||||
custom_timeout = httpx.Timeout(
|
||||
connect=120.0, read=upload_timeout, write=upload_timeout, pool=120.0,
|
||||
connect=120.0,
|
||||
read=upload_timeout,
|
||||
write=upload_timeout,
|
||||
pool=120.0,
|
||||
)
|
||||
|
||||
last_exc: Exception | None = None
|
||||
|
|
@ -135,12 +138,12 @@ async def parse_with_llamacloud(
|
|||
result = await parser.aparse(str(file_path))
|
||||
content = _extract_markdown(result).strip()
|
||||
if not content:
|
||||
raise LlamaCloudError(
|
||||
f"LlamaCloud returned empty content for {file_path}"
|
||||
)
|
||||
raise LlamaCloudError(f"LlamaCloud returned empty content for {file_path}")
|
||||
logger.info(
|
||||
"LlamaCloud OK: %s (%s) -> %d chars",
|
||||
file_path, parse_mode, len(content),
|
||||
file_path,
|
||||
parse_mode,
|
||||
len(content),
|
||||
)
|
||||
return content
|
||||
|
||||
|
|
@ -156,7 +159,10 @@ async def parse_with_llamacloud(
|
|||
sleep_for = delay + jitter
|
||||
logger.warning(
|
||||
"LlamaCloud attempt %d/%d failed (%s); retrying in %.1fs",
|
||||
attempt, _MAX_RETRIES, type(last_exc).__name__, sleep_for,
|
||||
attempt,
|
||||
_MAX_RETRIES,
|
||||
type(last_exc).__name__,
|
||||
sleep_for,
|
||||
)
|
||||
await asyncio.sleep(sleep_for)
|
||||
|
||||
|
|
|
|||
|
|
@ -116,11 +116,7 @@ def _normalise_paragraphs(text: str) -> list[str]:
|
|||
|
||||
|
||||
def _escape_html(text: str) -> str:
|
||||
return (
|
||||
text.replace("&", "&")
|
||||
.replace("<", "<")
|
||||
.replace(">", ">")
|
||||
)
|
||||
return text.replace("&", "&").replace("<", "<").replace(">", ">")
|
||||
|
||||
|
||||
def render_pdf(
|
||||
|
|
|
|||
|
|
@ -121,9 +121,7 @@ class OpenRouterPdfProvider:
|
|||
body: dict[str, Any] = {
|
||||
"model": self._model,
|
||||
"messages": messages,
|
||||
"plugins": [
|
||||
{"id": "file-parser", "pdf": {"engine": self._engine.value}}
|
||||
],
|
||||
"plugins": [{"id": "file-parser", "pdf": {"engine": self._engine.value}}],
|
||||
}
|
||||
if max_tokens:
|
||||
body["max_tokens"] = max_tokens
|
||||
|
|
|
|||
|
|
@ -177,7 +177,9 @@ class Benchmark(Protocol):
|
|||
def add_run_args(self, parser: argparse.ArgumentParser) -> None: # pragma: no cover - protocol
|
||||
"""Add benchmark-specific flags to ``run <suite> <benchmark>``."""
|
||||
|
||||
def report_section(self, artifacts: list[RunArtifact]) -> ReportSection: # pragma: no cover - protocol
|
||||
def report_section(
|
||||
self, artifacts: list[RunArtifact]
|
||||
) -> ReportSection: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
|
||||
|
|
@ -224,9 +226,7 @@ def get(suite: str, name: str) -> Benchmark:
|
|||
return _REGISTRY[(suite, name)]
|
||||
except KeyError as exc:
|
||||
available = ", ".join(f"{s}/{n}" for s, n in sorted(_REGISTRY)) or "<none>"
|
||||
raise KeyError(
|
||||
f"Unknown benchmark '{suite}/{name}'. Registered: {available}"
|
||||
) from exc
|
||||
raise KeyError(f"Unknown benchmark '{suite}/{name}'. Registered: {available}") from exc
|
||||
|
||||
|
||||
def list_suites() -> list[str]:
|
||||
|
|
|
|||
|
|
@ -45,10 +45,7 @@ def format_scenario_md(extra: Mapping[str, Any] | None) -> str:
|
|||
"(text-only model can't see images) — that's the point."
|
||||
)
|
||||
else:
|
||||
body = (
|
||||
f"- Scenario: head-to-head — both arms answer with `{surf_slug}` "
|
||||
"via OpenRouter."
|
||||
)
|
||||
body = f"- Scenario: head-to-head — both arms answer with `{surf_slug}` via OpenRouter."
|
||||
if vision_slug:
|
||||
body += f" SurfSense ingest VLM: `{vision_slug}`."
|
||||
|
||||
|
|
|
|||
|
|
@ -60,7 +60,5 @@ def discover_suites() -> list[str]:
|
|||
importlib.import_module(benchmark_name)
|
||||
imported.append(benchmark_name)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(
|
||||
"Failed to import benchmark %s: %s", benchmark_name, exc
|
||||
)
|
||||
logger.warning("Failed to import benchmark %s: %s", benchmark_name, exc)
|
||||
return imported
|
||||
|
|
|
|||
|
|
@ -154,13 +154,11 @@ async def run_ingest(
|
|||
if not batches:
|
||||
logger.warning("Discipline %s produced 0 batches; skipping upload", discipline)
|
||||
continue
|
||||
logger.info(
|
||||
"Uploading %d batches for discipline %s", len(batches), discipline
|
||||
)
|
||||
logger.info("Uploading %d batches for discipline %s", len(batches), discipline)
|
||||
upload_result = await docs_client.upload(
|
||||
files=[b.path for b in batches],
|
||||
search_space_id=ctx.search_space_id,
|
||||
use_vision_llm=settings.use_vision_llm,
|
||||
use_vision_llm=settings.use_vision_llm,
|
||||
processing_mode=settings.processing_mode,
|
||||
)
|
||||
new_doc_ids = list(upload_result.document_ids)
|
||||
|
|
@ -177,9 +175,7 @@ async def run_ingest(
|
|||
)
|
||||
title_to_doc = {s.title: s.document_id for s in statuses}
|
||||
|
||||
per_discipline_path = (
|
||||
ctx.maps_dir() / f"cure_corpus_map_{discipline}.jsonl"
|
||||
)
|
||||
per_discipline_path = ctx.maps_dir() / f"cure_corpus_map_{discipline}.jsonl"
|
||||
with per_discipline_path.open("w", encoding="utf-8") as fh:
|
||||
fh.write(settings_header_line(settings) + "\n")
|
||||
for batch in batches:
|
||||
|
|
@ -202,9 +198,7 @@ async def run_ingest(
|
|||
try:
|
||||
chunks = await docs_client.list_chunks(int(doc_id))
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(
|
||||
"Failed to list chunks for doc_id=%s: %s", doc_id, exc
|
||||
)
|
||||
logger.warning("Failed to list chunks for doc_id=%s: %s", doc_id, exc)
|
||||
continue
|
||||
for chunk in chunks:
|
||||
fh.write(
|
||||
|
|
|
|||
|
|
@ -191,12 +191,15 @@ class CureBenchmark:
|
|||
|
||||
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument("--lang", default="en", choices=("en", "es", "fr"))
|
||||
parser.add_argument("--discipline", default=None,
|
||||
help="Restrict to one discipline (default: all ingested).")
|
||||
parser.add_argument(
|
||||
"--discipline", default=None, help="Restrict to one discipline (default: all ingested)."
|
||||
)
|
||||
parser.add_argument("--n", dest="sample_n", type=int, default=None)
|
||||
parser.add_argument("--concurrency", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--max-passages-per-discipline", type=int, default=None,
|
||||
"--max-passages-per-discipline",
|
||||
type=int,
|
||||
default=None,
|
||||
help="(ingest only) cap corpus rows per discipline for smoke testing.",
|
||||
)
|
||||
# Per-upload knobs forwarded to /documents/fileupload at ingest;
|
||||
|
|
@ -233,11 +236,13 @@ class CureBenchmark:
|
|||
|
||||
# Disciplines to query are determined by the per-discipline maps
|
||||
# actually present (either user-filtered or whatever was ingested).
|
||||
ingested_disciplines = sorted({
|
||||
row_disc
|
||||
for path in maps_dir.glob("cure_corpus_map_*.jsonl")
|
||||
for row_disc in [path.stem[len("cure_corpus_map_"):]]
|
||||
})
|
||||
ingested_disciplines = sorted(
|
||||
{
|
||||
row_disc
|
||||
for path in maps_dir.glob("cure_corpus_map_*.jsonl")
|
||||
for row_disc in [path.stem[len("cure_corpus_map_") :]]
|
||||
}
|
||||
)
|
||||
if discipline_filter:
|
||||
disciplines = [discipline_filter]
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -55,15 +55,15 @@ def _hf_hub_download(*args, **kwargs):
|
|||
|
||||
@dataclass
|
||||
class MedXpertQuestion:
|
||||
qid: str # e.g. "MM-26"
|
||||
question: str # full question text (case + ask)
|
||||
options: dict[str, str] # A-E
|
||||
label: str # "A".."E"
|
||||
image_files: list[str] # filenames inside images.zip
|
||||
qid: str # e.g. "MM-26"
|
||||
question: str # full question text (case + ask)
|
||||
options: dict[str, str] # A-E
|
||||
label: str # "A".."E"
|
||||
image_files: list[str] # filenames inside images.zip
|
||||
medical_task: str
|
||||
body_system: str
|
||||
question_type: str
|
||||
split: str # "test" or "dev"
|
||||
split: str # "test" or "dev"
|
||||
|
||||
|
||||
def _load_jsonl(path: Path, *, split: str) -> list[MedXpertQuestion]:
|
||||
|
|
@ -84,17 +84,19 @@ def _load_jsonl(path: Path, *, split: str) -> list[MedXpertQuestion]:
|
|||
images = row.get("images") or []
|
||||
if not isinstance(images, list):
|
||||
images = []
|
||||
out.append(MedXpertQuestion(
|
||||
qid=qid,
|
||||
question=question,
|
||||
options=opts,
|
||||
label=label,
|
||||
image_files=[str(x).strip() for x in images if str(x).strip()],
|
||||
medical_task=str(row.get("medical_task") or "").strip(),
|
||||
body_system=str(row.get("body_system") or "").strip(),
|
||||
question_type=str(row.get("question_type") or "").strip(),
|
||||
split=split,
|
||||
))
|
||||
out.append(
|
||||
MedXpertQuestion(
|
||||
qid=qid,
|
||||
question=question,
|
||||
options=opts,
|
||||
label=label,
|
||||
image_files=[str(x).strip() for x in images if str(x).strip()],
|
||||
medical_task=str(row.get("medical_task") or "").strip(),
|
||||
body_system=str(row.get("body_system") or "").strip(),
|
||||
question_type=str(row.get("question_type") or "").strip(),
|
||||
split=split,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
|
|
@ -204,7 +206,7 @@ async def _upload_pdfs(
|
|||
name_to_id: dict[str, int] = {}
|
||||
pdf_list = list(pdf_paths)
|
||||
for batch_start in range(0, len(pdf_list), batch_size):
|
||||
batch = pdf_list[batch_start:batch_start + batch_size]
|
||||
batch = pdf_list[batch_start : batch_start + batch_size]
|
||||
result = await docs_client.upload(
|
||||
files=batch,
|
||||
search_space_id=ctx.search_space_id,
|
||||
|
|
@ -226,8 +228,10 @@ async def _upload_pdfs(
|
|||
name_to_id[s.title] = s.document_id
|
||||
logger.info(
|
||||
"Uploaded MedXpertQA batch %d-%d: %d new, %d duplicate",
|
||||
batch_start, batch_start + len(batch),
|
||||
len(result.document_ids), len(result.duplicate_document_ids),
|
||||
batch_start,
|
||||
batch_start + len(batch),
|
||||
len(result.document_ids),
|
||||
len(result.duplicate_document_ids),
|
||||
)
|
||||
return name_to_id
|
||||
|
||||
|
|
@ -310,9 +314,11 @@ async def run_ingest(
|
|||
# Materialise into bench_dir so the path is stable.
|
||||
try:
|
||||
from os import link as _link
|
||||
|
||||
_link(local_zip, images_zip_local)
|
||||
except OSError:
|
||||
from shutil import copy2
|
||||
|
||||
copy2(local_zip, images_zip_local)
|
||||
_ensure_images_extracted(images_zip_local, images_dir)
|
||||
|
||||
|
|
@ -354,17 +360,22 @@ async def run_ingest(
|
|||
questions_jsonl = bench_dir / "questions.jsonl"
|
||||
with questions_jsonl.open("w", encoding="utf-8") as fh:
|
||||
for q in questions:
|
||||
fh.write(json.dumps({
|
||||
"qid": q.qid,
|
||||
"question": q.question,
|
||||
"options": q.options,
|
||||
"label": q.label,
|
||||
"image_files": q.image_files,
|
||||
"medical_task": q.medical_task,
|
||||
"body_system": q.body_system,
|
||||
"question_type": q.question_type,
|
||||
"split": q.split,
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
"qid": q.qid,
|
||||
"question": q.question,
|
||||
"options": q.options,
|
||||
"label": q.label,
|
||||
"image_files": q.image_files,
|
||||
"medical_task": q.medical_task,
|
||||
"body_system": q.body_system,
|
||||
"question_type": q.question_type,
|
||||
"split": q.split,
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
logger.info("Wrote %d MedXpertQA questions to %s", len(questions), questions_jsonl)
|
||||
|
||||
map_path = ctx.maps_dir() / "medxpertqa_doc_map.jsonl"
|
||||
|
|
@ -376,13 +387,18 @@ async def run_ingest(
|
|||
local = pdf_paths.get(q.qid)
|
||||
if local is None:
|
||||
continue
|
||||
fh.write(json.dumps({
|
||||
"qid": q.qid,
|
||||
"document_id": name_to_id.get(local.name),
|
||||
"pdf_path": str(local),
|
||||
"n_images": len(q.image_files),
|
||||
"split": q.split,
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
"qid": q.qid,
|
||||
"document_id": name_to_id.get(local.name),
|
||||
"pdf_path": str(local),
|
||||
"n_images": len(q.image_files),
|
||||
"split": q.split,
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
logger.info("Wrote MedXpertQA doc map to %s", map_path)
|
||||
|
||||
new_state = ctx.suite_state
|
||||
|
|
|
|||
|
|
@ -129,19 +129,21 @@ def _load_questions(
|
|||
n_images = int(map_row.get("n_images", 0))
|
||||
if require_images and n_images <= 0:
|
||||
continue
|
||||
out.append(MXQuestion(
|
||||
qid=qid,
|
||||
question=str(row.get("question") or ""),
|
||||
options={str(k).upper(): str(v) for k, v in (row.get("options") or {}).items()},
|
||||
label=str(row.get("label") or "").strip().upper(),
|
||||
medical_task=str(row.get("medical_task") or "").strip(),
|
||||
body_system=str(row.get("body_system") or "").strip(),
|
||||
question_type=str(row.get("question_type") or "").strip(),
|
||||
split=str(row.get("split") or ""),
|
||||
n_images=n_images,
|
||||
pdf_path=Path(map_row["pdf_path"]),
|
||||
document_id=map_row.get("document_id"),
|
||||
))
|
||||
out.append(
|
||||
MXQuestion(
|
||||
qid=qid,
|
||||
question=str(row.get("question") or ""),
|
||||
options={str(k).upper(): str(v) for k, v in (row.get("options") or {}).items()},
|
||||
label=str(row.get("label") or "").strip().upper(),
|
||||
medical_task=str(row.get("medical_task") or "").strip(),
|
||||
body_system=str(row.get("body_system") or "").strip(),
|
||||
question_type=str(row.get("question_type") or "").strip(),
|
||||
split=str(row.get("split") or ""),
|
||||
n_images=n_images,
|
||||
pdf_path=Path(map_row["pdf_path"]),
|
||||
document_id=map_row.get("document_id"),
|
||||
)
|
||||
)
|
||||
out.sort(key=lambda q: (q.split, q.qid))
|
||||
if sample_n is not None and sample_n > 0:
|
||||
out = out[:sample_n]
|
||||
|
|
@ -182,51 +184,81 @@ class MedXpertQAMMBenchmark:
|
|||
|
||||
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--split", default="test", choices=["test", "dev", "all"],
|
||||
"--split",
|
||||
default="test",
|
||||
choices=["test", "dev", "all"],
|
||||
help="Which MedXpertQA-MM split to run (default: test).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--task", default="all",
|
||||
"--task",
|
||||
default="all",
|
||||
help="Filter by medical_task value (e.g. Diagnosis, Treatment, Basic Medicine).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--body-system", dest="body_filter", default="all",
|
||||
"--body-system",
|
||||
dest="body_filter",
|
||||
default="all",
|
||||
help="Filter by body_system value (e.g. Cardiovascular, Lymphatic).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--require-images", dest="require_images", action="store_true",
|
||||
"--require-images",
|
||||
dest="require_images",
|
||||
action="store_true",
|
||||
help="Skip rare MM rows that ended up with zero resolvable images.",
|
||||
)
|
||||
parser.add_argument("--n", dest="sample_n", type=int, default=None,
|
||||
help="Run only the first N questions after filters apply.")
|
||||
parser.add_argument("--concurrency", type=int, default=4,
|
||||
help="Parallel question workers per arm.")
|
||||
parser.add_argument("--no-mentions", dest="no_mentions", action="store_true",
|
||||
help="SurfSense arm: skip mentioned_document_ids (unscoped retrieval).")
|
||||
parser.add_argument(
|
||||
"--pdf-engine", default="native",
|
||||
"--n",
|
||||
dest="sample_n",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Run only the first N questions after filters apply.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, default=4, help="Parallel question workers per arm."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-mentions",
|
||||
dest="no_mentions",
|
||||
action="store_true",
|
||||
help="SurfSense arm: skip mentioned_document_ids (unscoped retrieval).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pdf-engine",
|
||||
default="native",
|
||||
choices=[e.value for e in PdfEngine],
|
||||
help="OpenRouter file-parser engine for the native arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-output-tokens", type=int, default=512,
|
||||
"--max-output-tokens",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Cap on completion length for both arms.",
|
||||
)
|
||||
# Ingest-only knobs (forwarded by the CLI to ingest.run_ingest).
|
||||
parser.add_argument(
|
||||
"--max-questions", dest="max_questions", type=int, default=None,
|
||||
"--max-questions",
|
||||
dest="max_questions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="(ingest only) cap on number of MM questions to render + upload.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload-batch-size", dest="upload_batch_size", type=int, default=8,
|
||||
"--upload-batch-size",
|
||||
dest="upload_batch_size",
|
||||
type=int,
|
||||
default=8,
|
||||
help="(ingest only) PDFs per fileupload call.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-upload", dest="skip_upload", action="store_true",
|
||||
"--skip-upload",
|
||||
dest="skip_upload",
|
||||
action="store_true",
|
||||
help="(ingest only) render PDFs locally but don't push to SurfSense.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--include-dev", dest="include_dev", action="store_true",
|
||||
"--include-dev",
|
||||
dest="include_dev",
|
||||
action="store_true",
|
||||
help="(ingest only) shorthand for --split all.",
|
||||
)
|
||||
# Per-upload knobs forwarded to /documents/fileupload at ingest;
|
||||
|
|
@ -270,7 +302,8 @@ class MedXpertQAMMBenchmark:
|
|||
|
||||
doc_map, ingest_settings = _load_doc_map(map_path)
|
||||
questions = _load_questions(
|
||||
questions_jsonl, doc_map,
|
||||
questions_jsonl,
|
||||
doc_map,
|
||||
split_filter=split_filter,
|
||||
task_filter=task_filter if task_filter != "all" else None,
|
||||
body_filter=body_filter if body_filter != "all" else None,
|
||||
|
|
@ -378,13 +411,18 @@ class MedXpertQAMMBenchmark:
|
|||
|
||||
manifest_path = run_dir / "run_artifact.json"
|
||||
manifest_path.write_text(
|
||||
json.dumps({
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
}, indent=2, sort_keys=True) + "\n",
|
||||
json.dumps(
|
||||
{
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
},
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return artifact
|
||||
|
|
@ -536,8 +574,12 @@ def _compute_metrics(
|
|||
cost_pct = _safe_pct(surf_cost_agg.mean, native_cost_agg.mean)
|
||||
lat_pct = _safe_pct(surf_lat_agg.median, native_lat_agg.median)
|
||||
|
||||
per_task = _per_field(questions, native_correct, surf_correct, key=lambda q: q.medical_task or "unknown")
|
||||
per_body = _per_field(questions, native_correct, surf_correct, key=lambda q: q.body_system or "unknown")
|
||||
per_task = _per_field(
|
||||
questions, native_correct, surf_correct, key=lambda q: q.medical_task or "unknown"
|
||||
)
|
||||
per_body = _per_field(
|
||||
questions, native_correct, surf_correct, key=lambda q: q.body_system or "unknown"
|
||||
)
|
||||
|
||||
return {
|
||||
"native": {
|
||||
|
|
@ -593,8 +635,7 @@ def _per_field(
|
|||
"native_accuracy": (sum(n_correct) / len(pairs)) if pairs else 0.0,
|
||||
"surfsense_accuracy": (sum(s_correct) / len(pairs)) if pairs else 0.0,
|
||||
"delta_accuracy_pp": (
|
||||
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs)
|
||||
if pairs else 0.0
|
||||
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs) if pairs else 0.0
|
||||
),
|
||||
}
|
||||
return out
|
||||
|
|
|
|||
|
|
@ -48,9 +48,7 @@ from ....core.registry import RunContext
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
MIRAGE_BENCHMARK_URL = (
|
||||
"https://raw.githubusercontent.com/Teddy-XiongGZ/MIRAGE/main/benchmark.json"
|
||||
)
|
||||
MIRAGE_BENCHMARK_URL = "https://raw.githubusercontent.com/Teddy-XiongGZ/MIRAGE/main/benchmark.json"
|
||||
# Upstream only ships ONE zip — top-10k retrievals across 5 retrievers,
|
||||
# ~16 GB. We default to skipping it (see `--skip-snippet-filter`) and
|
||||
# ingesting the chosen corpus in full; this URL is only fetched when
|
||||
|
|
@ -100,8 +98,7 @@ def _reuse_cached_dest(dest: Path, *, expect_zip: bool, label: str) -> Path | No
|
|||
return None
|
||||
if expect_zip and not _is_valid_zip(dest):
|
||||
logger.warning(
|
||||
"Cached %s at %s failed ZIP validation (size=%d B); deleting "
|
||||
"and re-downloading.",
|
||||
"Cached %s at %s failed ZIP validation (size=%d B); deleting and re-downloading.",
|
||||
label,
|
||||
dest,
|
||||
dest.stat().st_size,
|
||||
|
|
@ -176,10 +173,13 @@ async def _fetch_to_path(
|
|||
)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(timeout_s, connect=20.0),
|
||||
follow_redirects=True,
|
||||
) as client, client.stream("GET", url, headers=headers) as response:
|
||||
async with (
|
||||
httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(timeout_s, connect=20.0),
|
||||
follow_redirects=True,
|
||||
) as client,
|
||||
client.stream("GET", url, headers=headers) as response,
|
||||
):
|
||||
if existing_bytes and response.status_code == 200:
|
||||
logger.warning(
|
||||
"Server ignored Range header for %s; restarting from 0.",
|
||||
|
|
@ -223,7 +223,7 @@ async def _fetch_to_path(
|
|||
raise
|
||||
except _RETRYABLE_NET_EXC as exc:
|
||||
last_exc = exc
|
||||
wait = min(60.0, 2.0 ** attempt)
|
||||
wait = min(60.0, 2.0**attempt)
|
||||
logger.warning(
|
||||
"Network error fetching %s (%s: %s); retrying in %.0fs.",
|
||||
label,
|
||||
|
|
@ -236,7 +236,7 @@ async def _fetch_to_path(
|
|||
last_exc = exc
|
||||
# Truncated body — drop the partial and retry from scratch.
|
||||
partial.unlink(missing_ok=True)
|
||||
wait = min(60.0, 2.0 ** attempt)
|
||||
wait = min(60.0, 2.0**attempt)
|
||||
logger.warning(
|
||||
"Truncated ZIP for %s; restarting from byte 0 in %.0fs.",
|
||||
label,
|
||||
|
|
@ -278,9 +278,9 @@ class _LargeDownloadAbort(RuntimeError):
|
|||
"""Raised when a download exceeds the safety threshold without opt-in."""
|
||||
|
||||
def __init__(self, label: str, size_bytes: int) -> None:
|
||||
gb = size_bytes / (1024 ** 3)
|
||||
gb = size_bytes / (1024**3)
|
||||
super().__init__(
|
||||
f"{label} would download ~{gb:.1f} GB, above the {_LARGE_DOWNLOAD_BYTES / (1024 ** 3):.0f} GB safety cap. "
|
||||
f"{label} would download ~{gb:.1f} GB, above the {_LARGE_DOWNLOAD_BYTES / (1024**3):.0f} GB safety cap. "
|
||||
"Re-run with `--allow-large-download` to acknowledge, or use "
|
||||
"`--skip-snippet-filter` to bypass this download entirely and "
|
||||
"ingest the full corpus instead."
|
||||
|
|
@ -320,9 +320,7 @@ def _read_snippet_ids(zip_path: Path, *, tasks: list[str]) -> dict[str, set[str]
|
|||
return out
|
||||
|
||||
|
||||
def _load_corpus(
|
||||
corpus_name: str, snippet_ids: set[str] | None
|
||||
) -> Iterable[SnippetRow]:
|
||||
def _load_corpus(corpus_name: str, snippet_ids: set[str] | None) -> Iterable[SnippetRow]:
|
||||
"""Stream rows from a MedRAG HF corpus.
|
||||
|
||||
* ``snippet_ids=None`` → yield every row (full-corpus ingestion path).
|
||||
|
|
@ -541,10 +539,7 @@ async def run_ingest(
|
|||
logger.warning("Failed to list chunks for doc_id=%s: %s", doc_id, exc)
|
||||
continue
|
||||
for chunk in chunks:
|
||||
fh.write(
|
||||
json.dumps({"chunk_id": chunk.id, "document_id": doc_id})
|
||||
+ "\n"
|
||||
)
|
||||
fh.write(json.dumps({"chunk_id": chunk.id, "document_id": doc_id}) + "\n")
|
||||
|
||||
new_state = ctx.suite_state
|
||||
new_state.ingestion_maps["mirage"] = str(snippet_map_path)
|
||||
|
|
|
|||
|
|
@ -134,15 +134,23 @@ class MirageBenchmark:
|
|||
choices=("all", *_TASKS),
|
||||
help="Run a single task or all (default: all).",
|
||||
)
|
||||
parser.add_argument("--n", dest="sample_n", type=int, default=None,
|
||||
help="Stratified sample size across tasks.")
|
||||
parser.add_argument(
|
||||
"--n",
|
||||
dest="sample_n",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Stratified sample size across tasks.",
|
||||
)
|
||||
parser.add_argument("--concurrency", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--corpus", default="MedRAG/textbooks",
|
||||
"--corpus",
|
||||
default="MedRAG/textbooks",
|
||||
help="HF MedRAG corpus to ingest from (default: MedRAG/textbooks).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-snippets-per-task", type=int, default=None,
|
||||
"--max-snippets-per-task",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Cap the per-task ingestion to N snippets (smoke).",
|
||||
)
|
||||
# Mutually exclusive: by default we skip the upstream 16 GB
|
||||
|
|
@ -152,18 +160,24 @@ class MirageBenchmark:
|
|||
# --allow-large-download).
|
||||
snippet_group = parser.add_mutually_exclusive_group()
|
||||
snippet_group.add_argument(
|
||||
"--use-snippet-filter", dest="use_snippet_filter", action="store_true",
|
||||
"--use-snippet-filter",
|
||||
dest="use_snippet_filter",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Download retrieved_snippets_10k.zip (~16 GB) and "
|
||||
"filter the corpus to those ids before ingest. "
|
||||
"Default: skip and ingest entire corpus.",
|
||||
"filter the corpus to those ids before ingest. "
|
||||
"Default: skip and ingest entire corpus.",
|
||||
)
|
||||
snippet_group.add_argument(
|
||||
"--skip-snippet-filter", dest="use_snippet_filter", action="store_false",
|
||||
"--skip-snippet-filter",
|
||||
dest="use_snippet_filter",
|
||||
action="store_false",
|
||||
help="(Default) Skip the 16 GB upstream zip; ingest entire corpus.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--allow-large-download", action="store_true", default=False,
|
||||
"--allow-large-download",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Permit downloads larger than 2 GB (e.g. retrieved_snippets_10k.zip).",
|
||||
)
|
||||
# Per-upload knobs; ignored at run-time (runner reads the
|
||||
|
|
@ -196,16 +210,13 @@ class MirageBenchmark:
|
|||
"`python -m surfsense_evals ingest medical mirage` first."
|
||||
)
|
||||
benchmark = json.loads(bench_path.read_text(encoding="utf-8"))
|
||||
ingest_settings = read_settings_header(
|
||||
ctx.maps_dir() / "mirage_snippet_map.jsonl"
|
||||
)
|
||||
ingest_settings = read_settings_header(ctx.maps_dir() / "mirage_snippet_map.jsonl")
|
||||
questions = _load_questions(benchmark, tasks=tasks, sample_n=sample_n)
|
||||
if not questions:
|
||||
raise RuntimeError(
|
||||
f"No MIRAGE questions matched task={task_filter!r} sample_n={sample_n!r}."
|
||||
)
|
||||
logger.info("MIRAGE: scheduled %d questions across tasks %s",
|
||||
len(questions), tasks)
|
||||
logger.info("MIRAGE: scheduled %d questions across tasks %s", len(questions), tasks)
|
||||
|
||||
arm = SurfSenseArm(
|
||||
client=ctx.new_chat_client(),
|
||||
|
|
@ -255,7 +266,10 @@ class MirageBenchmark:
|
|||
per_task_acc[task] = acc.to_dict()
|
||||
|
||||
macro = macro_accuracy(
|
||||
{t: accuracy_with_wilson_ci(d["n_correct"], d["n_total"]) for t, d in per_task_acc.items()}
|
||||
{
|
||||
t: accuracy_with_wilson_ci(d["n_correct"], d["n_total"])
|
||||
for t, d in per_task_acc.items()
|
||||
}
|
||||
)
|
||||
metrics = {"per_task": per_task_acc, "macro_accuracy": macro}
|
||||
|
||||
|
|
|
|||
|
|
@ -112,8 +112,9 @@ def _grade_int(pred: str, gold: str) -> GradeResult:
|
|||
if p_match is None:
|
||||
return GradeResult(False, 0.0, "int_eq", str(p_match), str(g_val))
|
||||
p_val = int(p_match.group(0).replace(",", ""))
|
||||
return GradeResult(p_val == g_val, 1.0 if p_val == g_val else 0.0,
|
||||
"int_eq", str(p_val), str(g_val))
|
||||
return GradeResult(
|
||||
p_val == g_val, 1.0 if p_val == g_val else 0.0, "int_eq", str(p_val), str(g_val)
|
||||
)
|
||||
|
||||
|
||||
_FLOAT_RE = re.compile(r"-?\d+(?:[.,]\d+)?")
|
||||
|
|
@ -145,15 +146,15 @@ def _grade_list(pred: str, gold: str) -> GradeResult:
|
|||
return _grade_str(pred, gold)
|
||||
inter = g_items & p_items
|
||||
if not inter:
|
||||
return GradeResult(False, 0.0, "list_set",
|
||||
", ".join(sorted(p_items)),
|
||||
", ".join(sorted(g_items)))
|
||||
return GradeResult(
|
||||
False, 0.0, "list_set", ", ".join(sorted(p_items)), ", ".join(sorted(g_items))
|
||||
)
|
||||
precision = len(inter) / len(p_items) if p_items else 0.0
|
||||
recall = len(inter) / len(g_items)
|
||||
f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
|
||||
return GradeResult(f1 >= 0.999, f1, "list_set",
|
||||
", ".join(sorted(p_items)),
|
||||
", ".join(sorted(g_items)))
|
||||
return GradeResult(
|
||||
f1 >= 0.999, f1, "list_set", ", ".join(sorted(p_items)), ", ".join(sorted(g_items))
|
||||
)
|
||||
|
||||
|
||||
def _grade_none(pred: str, gold: str) -> GradeResult:
|
||||
|
|
@ -188,8 +189,11 @@ def _grade_none(pred: str, gold: str) -> GradeResult:
|
|||
expressed_unknown = True
|
||||
break
|
||||
return GradeResult(
|
||||
expressed_unknown, 1.0 if expressed_unknown else 0.0,
|
||||
"none_match", p, _normalise_text(gold),
|
||||
expressed_unknown,
|
||||
1.0 if expressed_unknown else 0.0,
|
||||
"none_match",
|
||||
p,
|
||||
_normalise_text(gold),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ logger = logging.getLogger(__name__)
|
|||
HF_REPO_ID = "yubo2333/MMLongBench-Doc"
|
||||
HF_REPO_TYPE = "dataset"
|
||||
|
||||
|
||||
# Lazy import: huggingface_hub + pyarrow are heavyweight; keep the
|
||||
# benchmark module importable on machines that have only the core
|
||||
# install (e.g. CI lint jobs).
|
||||
|
|
@ -63,11 +64,11 @@ def _list_repo_files() -> list[str]:
|
|||
|
||||
@dataclass
|
||||
class MMLongBenchQuestion:
|
||||
doc_id: str # filename inside the documents/ folder
|
||||
doc_id: str # filename inside the documents/ folder
|
||||
doc_type: str
|
||||
question: str
|
||||
answer: str
|
||||
answer_format: str # Str / Int / Float / List / None
|
||||
answer_format: str # Str / Int / Float / List / None
|
||||
evidence_pages: list[int]
|
||||
evidence_sources: list[str]
|
||||
|
||||
|
|
@ -161,7 +162,9 @@ def _download_questions_parquet(cache_dir: Path) -> Path:
|
|||
)
|
||||
parquet_paths.append(Path(local))
|
||||
logger.info("Cached MMLongBench parquet shard %s -> %s", rel, local)
|
||||
return parquet_paths[0] if len(parquet_paths) == 1 else _merge_parquets(parquet_paths, cache_dir)
|
||||
return (
|
||||
parquet_paths[0] if len(parquet_paths) == 1 else _merge_parquets(parquet_paths, cache_dir)
|
||||
)
|
||||
|
||||
|
||||
def _merge_parquets(paths: list[Path], cache_dir: Path) -> Path:
|
||||
|
|
@ -221,7 +224,7 @@ async def _upload_pdfs(
|
|||
name_to_id: dict[str, int] = {}
|
||||
pdf_list = list(pdf_paths)
|
||||
for batch_start in range(0, len(pdf_list), batch_size):
|
||||
batch = pdf_list[batch_start:batch_start + batch_size]
|
||||
batch = pdf_list[batch_start : batch_start + batch_size]
|
||||
result = await docs_client.upload(
|
||||
files=batch,
|
||||
search_space_id=ctx.search_space_id,
|
||||
|
|
@ -243,8 +246,10 @@ async def _upload_pdfs(
|
|||
name_to_id[s.title] = s.document_id
|
||||
logger.info(
|
||||
"Uploaded MMLongBench batch %d-%d: %d new, %d duplicate",
|
||||
batch_start, batch_start + len(batch),
|
||||
len(result.document_ids), len(result.duplicate_document_ids),
|
||||
batch_start,
|
||||
batch_start + len(batch),
|
||||
len(result.document_ids),
|
||||
len(result.duplicate_document_ids),
|
||||
)
|
||||
return name_to_id
|
||||
|
||||
|
|
@ -299,15 +304,20 @@ async def run_ingest(
|
|||
questions_jsonl = bench_dir / "questions.jsonl"
|
||||
with questions_jsonl.open("w", encoding="utf-8") as fh:
|
||||
for q in questions:
|
||||
fh.write(json.dumps({
|
||||
"doc_id": q.doc_id,
|
||||
"doc_type": q.doc_type,
|
||||
"question": q.question,
|
||||
"answer": q.answer,
|
||||
"answer_format": q.answer_format,
|
||||
"evidence_pages": q.evidence_pages,
|
||||
"evidence_sources": q.evidence_sources,
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
"doc_id": q.doc_id,
|
||||
"doc_type": q.doc_type,
|
||||
"question": q.question,
|
||||
"answer": q.answer,
|
||||
"answer_format": q.answer_format,
|
||||
"evidence_pages": q.evidence_pages,
|
||||
"evidence_sources": q.evidence_sources,
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
logger.info("Wrote %d MMLongBench questions to %s", len(questions), questions_jsonl)
|
||||
|
||||
# Step 2: download unique PDFs
|
||||
|
|
@ -348,12 +358,17 @@ async def run_ingest(
|
|||
local = pdf_paths.get(doc_id)
|
||||
if local is None:
|
||||
continue
|
||||
fh.write(json.dumps({
|
||||
"doc_id": doc_id,
|
||||
"document_id": name_to_id.get(local.name),
|
||||
"pdf_path": str(local),
|
||||
"n_questions": sum(1 for q in questions if q.doc_id == doc_id),
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"document_id": name_to_id.get(local.name),
|
||||
"pdf_path": str(local),
|
||||
"n_questions": sum(1 for q in questions if q.doc_id == doc_id),
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
logger.info("Wrote MMLongBench doc map to %s", map_path)
|
||||
|
||||
new_state = ctx.suite_state
|
||||
|
|
|
|||
|
|
@ -18,10 +18,7 @@ _FORMAT_HINTS: dict[str, str] = {
|
|||
"Respond with the answer as a short phrase, no full sentence. "
|
||||
"Format your final line as `Answer: <text>`."
|
||||
),
|
||||
"int": (
|
||||
"Respond with a single integer only. "
|
||||
"Format your final line as `Answer: <integer>`."
|
||||
),
|
||||
"int": ("Respond with a single integer only. Format your final line as `Answer: <integer>`."),
|
||||
"float": (
|
||||
"Respond with a single decimal number only (no units). "
|
||||
"Format your final line as `Answer: <number>`."
|
||||
|
|
|
|||
|
|
@ -58,8 +58,8 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
@dataclass
|
||||
class MMLBQuestion:
|
||||
qid: str # synthesised from doc_id + index
|
||||
doc_id: str # filename inside the documents/ folder
|
||||
qid: str # synthesised from doc_id + index
|
||||
doc_id: str # filename inside the documents/ folder
|
||||
doc_type: str
|
||||
question: str
|
||||
gold_answer: str
|
||||
|
|
@ -126,18 +126,20 @@ def _load_questions(
|
|||
continue
|
||||
idx = per_doc_counter.get(doc_id, 0)
|
||||
per_doc_counter[doc_id] = idx + 1
|
||||
out.append(MMLBQuestion(
|
||||
qid=f"{doc_id}::Q{idx:03d}",
|
||||
doc_id=doc_id,
|
||||
doc_type=str(row.get("doc_type") or "").strip(),
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=gold,
|
||||
answer_format=answer_format,
|
||||
evidence_pages=list(row.get("evidence_pages") or []),
|
||||
evidence_sources=list(row.get("evidence_sources") or []),
|
||||
pdf_path=Path(map_row["pdf_path"]),
|
||||
document_id=map_row.get("document_id"),
|
||||
))
|
||||
out.append(
|
||||
MMLBQuestion(
|
||||
qid=f"{doc_id}::Q{idx:03d}",
|
||||
doc_id=doc_id,
|
||||
doc_type=str(row.get("doc_type") or "").strip(),
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=gold,
|
||||
answer_format=answer_format,
|
||||
evidence_pages=list(row.get("evidence_pages") or []),
|
||||
evidence_sources=list(row.get("evidence_sources") or []),
|
||||
pdf_path=Path(map_row["pdf_path"]),
|
||||
document_id=map_row.get("document_id"),
|
||||
)
|
||||
)
|
||||
out.sort(key=lambda q: (q.doc_id, q.qid))
|
||||
if sample_n is not None and sample_n > 0:
|
||||
out = out[:sample_n]
|
||||
|
|
@ -202,41 +204,61 @@ class MMLongBenchDocBenchmark:
|
|||
help="Filter to one answer format. 'none' = unanswerable probes only.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n", dest="sample_n", type=int, default=None,
|
||||
"--n",
|
||||
dest="sample_n",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Run only the first N questions after filters apply.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-unanswerable", dest="skip_unanswerable", action="store_true",
|
||||
"--skip-unanswerable",
|
||||
dest="skip_unanswerable",
|
||||
action="store_true",
|
||||
help="Drop ~22%% unanswerable questions (use to compare against baselines that don't include them).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, default=4,
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Parallel question workers per arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-mentions", dest="no_mentions", action="store_true",
|
||||
"--no-mentions",
|
||||
dest="no_mentions",
|
||||
action="store_true",
|
||||
help="SurfSense arm: skip mentioned_document_ids (unscoped retrieval).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pdf-engine", default="native",
|
||||
"--pdf-engine",
|
||||
default="native",
|
||||
choices=[e.value for e in PdfEngine],
|
||||
help="OpenRouter file-parser engine for the native arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-output-tokens", type=int, default=512,
|
||||
"--max-output-tokens",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Cap on completion length for both arms.",
|
||||
)
|
||||
# Ingest-only knobs (forwarded by the CLI to ingest.run_ingest).
|
||||
parser.add_argument(
|
||||
"--max-docs", dest="max_docs", type=int, default=None,
|
||||
"--max-docs",
|
||||
dest="max_docs",
|
||||
type=int,
|
||||
default=None,
|
||||
help="(ingest only) cap on number of unique PDFs to download + upload.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload-batch-size", dest="upload_batch_size", type=int, default=8,
|
||||
"--upload-batch-size",
|
||||
dest="upload_batch_size",
|
||||
type=int,
|
||||
default=8,
|
||||
help="(ingest only) PDFs per fileupload call.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-upload", dest="skip_upload", action="store_true",
|
||||
"--skip-upload",
|
||||
dest="skip_upload",
|
||||
action="store_true",
|
||||
help="(ingest only) cache PDFs locally but don't push to SurfSense.",
|
||||
)
|
||||
# Per-upload knobs forwarded to /documents/fileupload at ingest;
|
||||
|
|
@ -278,7 +300,8 @@ class MMLongBenchDocBenchmark:
|
|||
|
||||
doc_map, ingest_settings = _load_doc_map(map_path)
|
||||
questions = _load_questions(
|
||||
questions_jsonl, doc_map,
|
||||
questions_jsonl,
|
||||
doc_map,
|
||||
doc_filter=doc_filter,
|
||||
format_filter=None if format_filter == "all" else format_filter,
|
||||
sample_n=sample_n,
|
||||
|
|
@ -292,9 +315,7 @@ class MMLongBenchDocBenchmark:
|
|||
|
||||
api_key = os.environ.get("OPENROUTER_API_KEY")
|
||||
if not api_key:
|
||||
raise RuntimeError(
|
||||
"OPENROUTER_API_KEY env var is required for the native arm."
|
||||
)
|
||||
raise RuntimeError("OPENROUTER_API_KEY env var is required for the native arm.")
|
||||
|
||||
# Native arm slug differs from SurfSense slug only in cost-arbitrage
|
||||
# scenario; otherwise both arms answer with provider_model.
|
||||
|
|
@ -362,18 +383,30 @@ class MMLongBenchDocBenchmark:
|
|||
"evidence_sources": q.evidence_sources,
|
||||
"document_id": q.document_id,
|
||||
}
|
||||
fh.write(json.dumps({
|
||||
**meta,
|
||||
**n_res.to_jsonl(),
|
||||
"graded": _grade_to_jsonl(n_g),
|
||||
}) + "\n")
|
||||
fh.write(json.dumps({
|
||||
**meta,
|
||||
**s_res.to_jsonl(),
|
||||
"graded": _grade_to_jsonl(s_g),
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
**meta,
|
||||
**n_res.to_jsonl(),
|
||||
"graded": _grade_to_jsonl(n_g),
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
**meta,
|
||||
**s_res.to_jsonl(),
|
||||
"graded": _grade_to_jsonl(s_g),
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
metrics = _compute_metrics(questions, native_results, surf_results, native_grades, surf_grades)
|
||||
metrics = _compute_metrics(
|
||||
questions, native_results, surf_results, native_grades, surf_grades
|
||||
)
|
||||
artifact = RunArtifact(
|
||||
suite=self.suite,
|
||||
benchmark=self.name,
|
||||
|
|
@ -398,13 +431,18 @@ class MMLongBenchDocBenchmark:
|
|||
|
||||
manifest_path = run_dir / "run_artifact.json"
|
||||
manifest_path.write_text(
|
||||
json.dumps({
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
}, indent=2, sort_keys=True) + "\n",
|
||||
json.dumps(
|
||||
{
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
},
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return artifact
|
||||
|
|
@ -450,9 +488,7 @@ class MMLongBenchDocBenchmark:
|
|||
f"(McNemar p={_fmt(delta.get('mcnemar_p_value'), 4)}, "
|
||||
f"method={delta.get('mcnemar_method')})"
|
||||
)
|
||||
body_lines.append(
|
||||
f" - F1 (mean): SurfSense {_pp(delta.get('f1_pp'))} pp"
|
||||
)
|
||||
body_lines.append(f" - F1 (mean): SurfSense {_pp(delta.get('f1_pp'))} pp")
|
||||
body_lines.append(
|
||||
f" - Bootstrap 95% CI on accuracy delta: "
|
||||
f"[{_pp(delta.get('bootstrap_ci_low'))}pp, {_pp(delta.get('bootstrap_ci_high'))}pp]"
|
||||
|
|
@ -472,8 +508,8 @@ class MMLongBenchDocBenchmark:
|
|||
for fmt, vals in sorted(per_format.items()):
|
||||
body_lines.append(
|
||||
f" - {fmt}: SurfSense {_pp(vals.get('delta_accuracy_pp'))} pp "
|
||||
f"(n={vals.get('n')}, native acc={vals.get('native_accuracy', 0)*100:.1f}%, "
|
||||
f"surf acc={vals.get('surfsense_accuracy', 0)*100:.1f}%)"
|
||||
f"(n={vals.get('n')}, native acc={vals.get('native_accuracy', 0) * 100:.1f}%, "
|
||||
f"surf acc={vals.get('surfsense_accuracy', 0) * 100:.1f}%)"
|
||||
)
|
||||
|
||||
return ReportSection(
|
||||
|
|
@ -576,8 +612,7 @@ def _compute_metrics(
|
|||
"native_accuracy": (sum(n_correct) / len(pairs)) if pairs else 0.0,
|
||||
"surfsense_accuracy": (sum(s_correct) / len(pairs)) if pairs else 0.0,
|
||||
"delta_accuracy_pp": (
|
||||
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs)
|
||||
if pairs else 0.0
|
||||
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs) if pairs else 0.0
|
||||
),
|
||||
}
|
||||
|
||||
|
|
@ -593,8 +628,12 @@ def _compute_metrics(
|
|||
"latency_ms_mean": native_latency_agg.mean,
|
||||
"latency_ms_median": native_latency_agg.median,
|
||||
"latency_ms_p95": native_latency_agg.p95,
|
||||
"input_tokens_mean": (sum(native_in_tokens) / len(native_in_tokens)) if native_in_tokens else 0.0,
|
||||
"output_tokens_mean": (sum(native_out_tokens) / len(native_out_tokens)) if native_out_tokens else 0.0,
|
||||
"input_tokens_mean": (sum(native_in_tokens) / len(native_in_tokens))
|
||||
if native_in_tokens
|
||||
else 0.0,
|
||||
"output_tokens_mean": (sum(native_out_tokens) / len(native_out_tokens))
|
||||
if native_out_tokens
|
||||
else 0.0,
|
||||
},
|
||||
"surfsense": {
|
||||
**surf_acc.to_dict(),
|
||||
|
|
|
|||
|
|
@ -53,9 +53,9 @@ logger = logging.getLogger(__name__)
|
|||
# Order matters for the manifest only (deterministic JSONL diffs);
|
||||
# the runner doesn't rely on it.
|
||||
PARSER_ARMS: tuple[tuple[str, str, str], ...] = (
|
||||
("azure_basic_lc", "azure", "basic"),
|
||||
("azure_premium_lc", "azure", "premium"),
|
||||
("llamacloud_basic_lc", "llamacloud", "basic"),
|
||||
("azure_basic_lc", "azure", "basic"),
|
||||
("azure_premium_lc", "azure", "premium"),
|
||||
("llamacloud_basic_lc", "llamacloud", "basic"),
|
||||
("llamacloud_premium_lc", "llamacloud", "premium"),
|
||||
)
|
||||
|
||||
|
|
@ -98,9 +98,7 @@ class PdfManifestRow:
|
|||
"pdf_path": str(self.pdf_path),
|
||||
"document_id": self.document_id,
|
||||
"pages": self.pages,
|
||||
"extractions": {
|
||||
arm: ext.to_jsonl() for arm, ext in self.extractions.items()
|
||||
},
|
||||
"extractions": {arm: ext.to_jsonl() for arm, ext in self.extractions.items()},
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -124,7 +122,9 @@ async def _run_one_extraction(
|
|||
markdown = await parse_with_azure_di(pdf_path, processing_mode=mode)
|
||||
elif parser == "llamacloud":
|
||||
markdown = await parse_with_llamacloud(
|
||||
pdf_path, processing_mode=mode, estimated_pages=estimated_pages,
|
||||
pdf_path,
|
||||
processing_mode=mode,
|
||||
estimated_pages=estimated_pages,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown parser {parser!r}")
|
||||
|
|
@ -168,14 +168,17 @@ async def _extract_one_pdf(
|
|||
error="(cached)",
|
||||
)
|
||||
logger.info(
|
||||
"Cached extraction reused: %s (%d chars)", out_path.name, len(cached),
|
||||
"Cached extraction reused: %s (%d chars)",
|
||||
out_path.name,
|
||||
len(cached),
|
||||
)
|
||||
coros.append(_noop())
|
||||
else:
|
||||
coros.append(
|
||||
_run_one_extraction(
|
||||
pdf_path,
|
||||
parser=parser, mode=mode,
|
||||
parser=parser,
|
||||
mode=mode,
|
||||
out_path=out_path,
|
||||
estimated_pages=estimated_pages,
|
||||
)
|
||||
|
|
@ -190,16 +193,24 @@ async def _extract_one_pdf(
|
|||
err_msg = f"{type(err).__name__}: {err}"
|
||||
logger.warning(
|
||||
"Extraction FAILED for %s [%s/%s]: %s",
|
||||
pdf_path.name, parser, mode, err_msg,
|
||||
pdf_path.name,
|
||||
parser,
|
||||
mode,
|
||||
err_msg,
|
||||
)
|
||||
out[arm_name] = ExtractionResult(
|
||||
arm=arm_name, parser=parser, mode=mode,
|
||||
status="failed", error=err_msg,
|
||||
arm=arm_name,
|
||||
parser=parser,
|
||||
mode=mode,
|
||||
status="failed",
|
||||
error=err_msg,
|
||||
)
|
||||
else:
|
||||
markdown, elapsed = result
|
||||
out[arm_name] = ExtractionResult(
|
||||
arm=arm_name, parser=parser, mode=mode,
|
||||
arm=arm_name,
|
||||
parser=parser,
|
||||
mode=mode,
|
||||
markdown_path=out_path,
|
||||
chars=len(markdown),
|
||||
elapsed_s=elapsed,
|
||||
|
|
@ -288,9 +299,7 @@ async def run_ingest(
|
|||
rows_in_scope = rows_in_scope[:max_docs]
|
||||
|
||||
if not rows_in_scope:
|
||||
raise RuntimeError(
|
||||
"No PDFs in scope for parser_compare. Check --docs / --max-docs."
|
||||
)
|
||||
raise RuntimeError("No PDFs in scope for parser_compare. Check --docs / --max-docs.")
|
||||
|
||||
bench_dir = ctx.benchmark_data_dir()
|
||||
extractions_dir = bench_dir / "extractions"
|
||||
|
|
@ -317,7 +326,8 @@ async def run_ingest(
|
|||
|
||||
logger.info(
|
||||
"parser_compare: extracting %d PDFs x 4 parsers (concurrency=%d)",
|
||||
len(rows_in_scope), pdf_concurrency,
|
||||
len(rows_in_scope),
|
||||
pdf_concurrency,
|
||||
)
|
||||
manifest_rows = await asyncio.gather(*(_process(r) for r in rows_in_scope))
|
||||
|
||||
|
|
@ -337,12 +347,13 @@ async def run_ingest(
|
|||
# Quick summary log
|
||||
total_extractions = sum(len(mr.extractions) for mr in manifest_rows)
|
||||
failures = sum(
|
||||
1 for mr in manifest_rows for ext in mr.extractions.values()
|
||||
if ext.status != "ok"
|
||||
1 for mr in manifest_rows for ext in mr.extractions.values() if ext.status != "ok"
|
||||
)
|
||||
logger.info(
|
||||
"parser_compare ingest done: %d PDFs, %d extractions, %d failures",
|
||||
len(manifest_rows), total_extractions, failures,
|
||||
len(manifest_rows),
|
||||
total_extractions,
|
||||
failures,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -34,10 +34,7 @@ _FORMAT_HINTS: dict[str, str] = {
|
|||
"Respond with the answer as a short phrase, no full sentence. "
|
||||
"Format your final line as `Answer: <text>`."
|
||||
),
|
||||
"int": (
|
||||
"Respond with a single integer only. "
|
||||
"Format your final line as `Answer: <integer>`."
|
||||
),
|
||||
"int": ("Respond with a single integer only. Format your final line as `Answer: <integer>`."),
|
||||
"float": (
|
||||
"Respond with a single decimal number only (no units). "
|
||||
"Format your final line as `Answer: <number>`."
|
||||
|
|
@ -69,11 +66,7 @@ _BASE_INSTRUCTION = (
|
|||
def build_native_pdf_prompt(question: str, *, answer_format: str) -> str:
|
||||
"""Prompt for ``NativePdfArm`` — PDF attached separately as a file part."""
|
||||
|
||||
return (
|
||||
f"{_BASE_INSTRUCTION}\n\n"
|
||||
f"Question: {question.strip()}\n\n"
|
||||
f"{_format_hint(answer_format)}\n"
|
||||
)
|
||||
return f"{_BASE_INSTRUCTION}\n\nQuestion: {question.strip()}\n\n{_format_hint(answer_format)}\n"
|
||||
|
||||
|
||||
def build_surfsense_prompt(question: str, *, answer_format: str) -> str:
|
||||
|
|
@ -82,11 +75,7 @@ def build_surfsense_prompt(question: str, *, answer_format: str) -> str:
|
|||
# SurfSense's agent already injects retrieved chunks via its tool
|
||||
# loop; the prompt only carries the user-visible question + format
|
||||
# hint, mirroring how a human asks the SurfSense UI.
|
||||
return (
|
||||
f"{_BASE_INSTRUCTION}\n\n"
|
||||
f"Question: {question.strip()}\n\n"
|
||||
f"{_format_hint(answer_format)}\n"
|
||||
)
|
||||
return f"{_BASE_INSTRUCTION}\n\nQuestion: {question.strip()}\n\n{_format_hint(answer_format)}\n"
|
||||
|
||||
|
||||
def build_long_context_prompt(
|
||||
|
|
@ -105,7 +94,7 @@ def build_long_context_prompt(
|
|||
|
||||
return (
|
||||
f"{_BASE_INSTRUCTION}\n\n"
|
||||
f"<document name=\"{document_label}\">\n"
|
||||
f'<document name="{document_label}">\n'
|
||||
f"{document_markdown.strip()}\n"
|
||||
f"</document>\n\n"
|
||||
f"Question: {question.strip()}\n\n"
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ logger = logging.getLogger(__name__)
|
|||
# Cost tariff (per the user's spec: $1 / 1k pages basic, $10 / 1k pages premium).
|
||||
# Held as dollars-per-page so per-PDF math is a pure multiply.
|
||||
PREPROCESS_USD_PER_PAGE = {
|
||||
"basic": 1.0 / 1000.0,
|
||||
"basic": 1.0 / 1000.0,
|
||||
"premium": 10.0 / 1000.0,
|
||||
}
|
||||
|
||||
|
|
@ -183,17 +183,19 @@ def _select_questions(
|
|||
if ext_blob.get("status") == "ok" and ext_blob.get("markdown_path"):
|
||||
extractions[arm_name] = Path(ext_blob["markdown_path"])
|
||||
|
||||
out.append(PCQuestion(
|
||||
qid=f"{doc_id}::Q{idx:03d}",
|
||||
doc_id=doc_id,
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=str(row.get("answer") or "").strip(),
|
||||
answer_format=answer_format,
|
||||
pdf_path=Path(map_row["pdf_path"]),
|
||||
document_id=map_row.get("document_id"),
|
||||
pages=int(map_row.get("pages", 0)),
|
||||
extractions=extractions,
|
||||
))
|
||||
out.append(
|
||||
PCQuestion(
|
||||
qid=f"{doc_id}::Q{idx:03d}",
|
||||
doc_id=doc_id,
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=str(row.get("answer") or "").strip(),
|
||||
answer_format=answer_format,
|
||||
pdf_path=Path(map_row["pdf_path"]),
|
||||
document_id=map_row.get("document_id"),
|
||||
pages=int(map_row.get("pages", 0)),
|
||||
extractions=extractions,
|
||||
)
|
||||
)
|
||||
per_doc_taken[doc_id] = per_doc_taken.get(doc_id, 0) + 1
|
||||
|
||||
out.sort(key=lambda q: (q.doc_id, q.qid))
|
||||
|
|
@ -242,65 +244,86 @@ class ParserCompareBenchmark:
|
|||
|
||||
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--docs", default=None,
|
||||
"--docs",
|
||||
default=None,
|
||||
help="Comma-separated doc_ids to include (default: all in manifest).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-per-doc", type=int, default=1,
|
||||
"--sample-per-doc",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Take the first N answerable questions per PDF (default 1).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-unanswerable", dest="skip_unanswerable",
|
||||
action="store_true", default=True,
|
||||
"--skip-unanswerable",
|
||||
dest="skip_unanswerable",
|
||||
action="store_true",
|
||||
default=True,
|
||||
help="Drop 'None' format probes (default true; we want signal not "
|
||||
"hallucination probes for n=5).",
|
||||
"hallucination probes for n=5).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--include-unanswerable", dest="skip_unanswerable",
|
||||
"--include-unanswerable",
|
||||
dest="skip_unanswerable",
|
||||
action="store_false",
|
||||
help="Override --skip-unanswerable; include unanswerable probes too.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-format", default=None,
|
||||
"--skip-format",
|
||||
default=None,
|
||||
help="Comma-separated answer_format values to skip (e.g. 'none,float').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, default=2,
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Parallel question workers per arm (default 2).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-mentions", dest="no_mentions", action="store_true",
|
||||
"--no-mentions",
|
||||
dest="no_mentions",
|
||||
action="store_true",
|
||||
help="SurfSense arm: skip mentioned_document_ids (full-corpus retrieval).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pdf-engine", default="native",
|
||||
"--pdf-engine",
|
||||
default="native",
|
||||
choices=[e.value for e in PdfEngine],
|
||||
help="OpenRouter file-parser engine for native_pdf arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-output-tokens", type=int, default=512,
|
||||
"--max-output-tokens",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Cap on completion length for every arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--llm-model", default="anthropic/claude-sonnet-4.5",
|
||||
"--llm-model",
|
||||
default="anthropic/claude-sonnet-4.5",
|
||||
help="OpenRouter slug used by the 5 OpenRouter-driven arms. "
|
||||
"SurfSense arm uses whatever provider_model is pinned on the suite.",
|
||||
"SurfSense arm uses whatever provider_model is pinned on the suite.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-arms", default=None,
|
||||
"--skip-arms",
|
||||
default=None,
|
||||
help="Comma-separated arm names to skip (e.g. 'llamacloud_premium_lc').",
|
||||
)
|
||||
# Ingest-only flags (forwarded by the CLI to ingest.run_ingest).
|
||||
parser.add_argument(
|
||||
"--max-docs", type=int, default=None,
|
||||
"--max-docs",
|
||||
type=int,
|
||||
default=None,
|
||||
help="(ingest only) cap number of unique PDFs to process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force-reextract", action="store_true",
|
||||
"--force-reextract",
|
||||
action="store_true",
|
||||
help="(ingest only) re-call parsers even if cached .md exists.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pdf-concurrency", type=int, default=2,
|
||||
"--pdf-concurrency",
|
||||
type=int,
|
||||
default=2,
|
||||
help="(ingest only) parallel PDFs (each fans out to 4 parsers).",
|
||||
)
|
||||
|
||||
|
|
@ -312,9 +335,7 @@ class ParserCompareBenchmark:
|
|||
from .ingest import run_ingest
|
||||
|
||||
docs_raw: str | None = opts.get("docs")
|
||||
docs_filter = (
|
||||
[d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
|
||||
)
|
||||
docs_filter = [d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
|
||||
await run_ingest(
|
||||
ctx,
|
||||
docs_filter=docs_filter,
|
||||
|
|
@ -329,15 +350,14 @@ class ParserCompareBenchmark:
|
|||
|
||||
async def run(self, ctx: RunContext, **opts: Any) -> RunArtifact:
|
||||
docs_raw: str | None = opts.get("docs")
|
||||
docs_filter = (
|
||||
[d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
|
||||
)
|
||||
docs_filter = [d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
|
||||
sample_per_doc = int(opts.get("sample_per_doc") or 1)
|
||||
skip_unanswerable = bool(opts.get("skip_unanswerable", True))
|
||||
skip_format_raw: str | None = opts.get("skip_format")
|
||||
skip_format = (
|
||||
[f.strip() for f in skip_format_raw.split(",") if f.strip()]
|
||||
if skip_format_raw else None
|
||||
if skip_format_raw
|
||||
else None
|
||||
)
|
||||
concurrency = int(opts.get("concurrency") or 2)
|
||||
no_mentions = bool(opts.get("no_mentions"))
|
||||
|
|
@ -346,8 +366,7 @@ class ParserCompareBenchmark:
|
|||
llm_model = str(opts.get("llm_model") or "anthropic/claude-sonnet-4.5")
|
||||
skip_arms_raw: str | None = opts.get("skip_arms")
|
||||
skip_arms = (
|
||||
{a.strip() for a in skip_arms_raw.split(",") if a.strip()}
|
||||
if skip_arms_raw else set()
|
||||
{a.strip() for a in skip_arms_raw.split(",") if a.strip()} if skip_arms_raw else set()
|
||||
)
|
||||
|
||||
active_arms = [a for a in ARM_NAMES if a not in skip_arms]
|
||||
|
|
@ -373,19 +392,20 @@ class ParserCompareBenchmark:
|
|||
|
||||
doc_map = _read_doc_map(map_path)
|
||||
questions = _select_questions(
|
||||
questions_jsonl, doc_map,
|
||||
questions_jsonl,
|
||||
doc_map,
|
||||
docs_filter=docs_filter,
|
||||
sample_per_doc=sample_per_doc,
|
||||
skip_unanswerable=skip_unanswerable,
|
||||
skip_format=skip_format,
|
||||
)
|
||||
if not questions:
|
||||
raise RuntimeError(
|
||||
"No questions matched filters; broaden --docs / --skip-format."
|
||||
)
|
||||
raise RuntimeError("No questions matched filters; broaden --docs / --skip-format.")
|
||||
logger.info(
|
||||
"parser_compare: scheduled %d questions across %d arms (%s)",
|
||||
len(questions), len(active_arms), ",".join(active_arms),
|
||||
len(questions),
|
||||
len(active_arms),
|
||||
",".join(active_arms),
|
||||
)
|
||||
|
||||
api_key = os.environ.get("OPENROUTER_API_KEY")
|
||||
|
|
@ -396,16 +416,20 @@ class ParserCompareBenchmark:
|
|||
arms: dict[str, Any] = {}
|
||||
if "native_pdf" in active_arms:
|
||||
native_provider = OpenRouterPdfProvider(
|
||||
api_key=api_key, base_url=ctx.config.openrouter_base_url,
|
||||
model=llm_model, engine=PdfEngine(pdf_engine_name),
|
||||
api_key=api_key,
|
||||
base_url=ctx.config.openrouter_base_url,
|
||||
model=llm_model,
|
||||
engine=PdfEngine(pdf_engine_name),
|
||||
)
|
||||
arms["native_pdf"] = NativePdfArm(
|
||||
provider=native_provider, max_output_tokens=max_output_tokens,
|
||||
provider=native_provider,
|
||||
max_output_tokens=max_output_tokens,
|
||||
)
|
||||
for arm_name, _, _ in PARSER_ARMS:
|
||||
if arm_name in active_arms:
|
||||
lc_provider = OpenRouterChatProvider(
|
||||
api_key=api_key, base_url=ctx.config.openrouter_base_url,
|
||||
api_key=api_key,
|
||||
base_url=ctx.config.openrouter_base_url,
|
||||
model=llm_model,
|
||||
)
|
||||
arms[arm_name] = BareLlmArm(
|
||||
|
|
@ -441,9 +465,7 @@ class ParserCompareBenchmark:
|
|||
def _lc_req(q: PCQuestion, arm_name: str) -> ArmRequest:
|
||||
md_path = q.extractions.get(arm_name)
|
||||
if md_path is None or not md_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Missing extraction for {arm_name} on {q.doc_id}"
|
||||
)
|
||||
raise FileNotFoundError(f"Missing extraction for {arm_name} on {q.doc_id}")
|
||||
markdown = md_path.read_text(encoding="utf-8")
|
||||
return ArmRequest(
|
||||
question_id=q.qid,
|
||||
|
|
@ -483,14 +505,15 @@ class ParserCompareBenchmark:
|
|||
|
||||
# Run all arms in parallel (each arm bounded by `concurrency`).
|
||||
per_arm_tasks: dict[str, list] = {
|
||||
arm_name: [_answer_one(arm_name, q) for q in questions]
|
||||
for arm_name in active_arms
|
||||
arm_name: [_answer_one(arm_name, q) for q in questions] for arm_name in active_arms
|
||||
}
|
||||
per_arm_results: dict[str, list[ArmResult]] = {}
|
||||
gathered = await asyncio.gather(*[
|
||||
_gather_with_limit(per_arm_tasks[arm_name], concurrency=concurrency)
|
||||
for arm_name in active_arms
|
||||
])
|
||||
gathered = await asyncio.gather(
|
||||
*[
|
||||
_gather_with_limit(per_arm_tasks[arm_name], concurrency=concurrency)
|
||||
for arm_name in active_arms
|
||||
]
|
||||
)
|
||||
for arm_name, results in zip(active_arms, gathered, strict=True):
|
||||
per_arm_results[arm_name] = results
|
||||
|
||||
|
|
@ -520,21 +543,29 @@ class ParserCompareBenchmark:
|
|||
for arm_name in active_arms:
|
||||
res = per_arm_results[arm_name][i]
|
||||
g = per_arm_grades[arm_name][i]
|
||||
fh.write(json.dumps({
|
||||
**base,
|
||||
**res.to_jsonl(),
|
||||
"graded": {
|
||||
"correct": g.correct,
|
||||
"f1": g.f1,
|
||||
"method": g.method,
|
||||
"normalised_pred": g.normalised_pred,
|
||||
"normalised_gold": g.normalised_gold,
|
||||
},
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
**base,
|
||||
**res.to_jsonl(),
|
||||
"graded": {
|
||||
"correct": g.correct,
|
||||
"f1": g.f1,
|
||||
"method": g.method,
|
||||
"normalised_pred": g.normalised_pred,
|
||||
"normalised_gold": g.normalised_gold,
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
# Aggregate per-arm metrics + cost
|
||||
metrics = _compute_metrics(
|
||||
questions, per_arm_results, per_arm_grades, active_arms,
|
||||
questions,
|
||||
per_arm_results,
|
||||
per_arm_grades,
|
||||
active_arms,
|
||||
)
|
||||
|
||||
artifact = RunArtifact(
|
||||
|
|
@ -564,13 +595,18 @@ class ParserCompareBenchmark:
|
|||
|
||||
manifest_path = run_dir / "run_artifact.json"
|
||||
manifest_path.write_text(
|
||||
json.dumps({
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
}, indent=2, sort_keys=True) + "\n",
|
||||
json.dumps(
|
||||
{
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
},
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return artifact
|
||||
|
|
@ -602,10 +638,7 @@ class ParserCompareBenchmark:
|
|||
f"(LLM: `{extra.get('llm_model', '?')}`, "
|
||||
f"engine: `{extra.get('pdf_engine', 'native')}`)."
|
||||
)
|
||||
body.append(
|
||||
"- Preprocess tariff: basic = $1 / 1k pages, "
|
||||
"premium = $10 / 1k pages."
|
||||
)
|
||||
body.append("- Preprocess tariff: basic = $1 / 1k pages, premium = $10 / 1k pages.")
|
||||
body.append("")
|
||||
body.append("### Per-arm summary")
|
||||
body.append("")
|
||||
|
|
@ -620,13 +653,13 @@ class ParserCompareBenchmark:
|
|||
continue
|
||||
body.append(
|
||||
f"| `{arm_name}` "
|
||||
f"| {row['accuracy']*100:.1f}% "
|
||||
f"| {row['accuracy'] * 100:.1f}% "
|
||||
f"({row['n_correct']}/{row['n']}) "
|
||||
f"| {row['f1_mean']*100:.1f}% "
|
||||
f"| {row['f1_mean'] * 100:.1f}% "
|
||||
f"| ${row['llm_cost_per_q']:.4f} "
|
||||
f"| ${row['preprocess_cost_total']:.4f} "
|
||||
f"| ${row['total_cost_per_q']:.4f} "
|
||||
f"| {row['latency_ms_median']/1000:.1f}s |"
|
||||
f"| {row['latency_ms_median'] / 1000:.1f}s |"
|
||||
)
|
||||
body.append("")
|
||||
|
||||
|
|
@ -679,8 +712,7 @@ class ParserCompareBenchmark:
|
|||
else:
|
||||
row_cells.append("✓" if g.get("correct") else "✗")
|
||||
body.append(
|
||||
f"| `{doc_id}` | {info.get('pages', '?')} | "
|
||||
+ " | ".join(row_cells) + " |"
|
||||
f"| `{doc_id}` | {info.get('pages', '?')} | " + " | ".join(row_cells) + " |"
|
||||
)
|
||||
|
||||
return ReportSection(
|
||||
|
|
@ -740,16 +772,16 @@ def _compute_metrics(
|
|||
preprocess_per_page = 0.0
|
||||
preprocess_label = "unknown"
|
||||
|
||||
preprocess_cost_total = sum(
|
||||
pages * preprocess_per_page for pages in pdf_pages.values()
|
||||
)
|
||||
preprocess_cost_total = sum(pages * preprocess_per_page for pages in pdf_pages.values())
|
||||
preprocess_cost_per_q = preprocess_cost_total / n if n else 0.0
|
||||
total_cost_per_q = llm_cost_per_q + preprocess_cost_per_q
|
||||
|
||||
latencies = sorted(int(r.latency_ms or 0) for r in results)
|
||||
latency_median = latencies[len(latencies) // 2] if latencies else 0
|
||||
latency_p95 = latencies[int(len(latencies) * 0.95)] if len(latencies) >= 20 else (
|
||||
latencies[-1] if latencies else 0
|
||||
latency_p95 = (
|
||||
latencies[int(len(latencies) * 0.95)]
|
||||
if len(latencies) >= 20
|
||||
else (latencies[-1] if latencies else 0)
|
||||
)
|
||||
|
||||
in_tokens = [int(r.input_tokens or 0) for r in results]
|
||||
|
|
@ -775,15 +807,21 @@ def _compute_metrics(
|
|||
# Per-PDF breakdown (correct / not for each arm)
|
||||
per_pdf: dict[str, dict[str, Any]] = {}
|
||||
for i, q in enumerate(questions):
|
||||
slot = per_pdf.setdefault(q.doc_id, {
|
||||
"pages": q.pages,
|
||||
"arms": {},
|
||||
})
|
||||
slot = per_pdf.setdefault(
|
||||
q.doc_id,
|
||||
{
|
||||
"pages": q.pages,
|
||||
"arms": {},
|
||||
},
|
||||
)
|
||||
for arm_name in active_arms:
|
||||
slot["arms"].setdefault(arm_name, {
|
||||
"correct": per_arm_grades[arm_name][i].correct,
|
||||
"f1": per_arm_grades[arm_name][i].f1,
|
||||
})
|
||||
slot["arms"].setdefault(
|
||||
arm_name,
|
||||
{
|
||||
"correct": per_arm_grades[arm_name][i].correct,
|
||||
"f1": per_arm_grades[arm_name][i].f1,
|
||||
},
|
||||
)
|
||||
|
||||
return {
|
||||
"per_arm": per_arm,
|
||||
|
|
|
|||
|
|
@ -80,7 +80,7 @@ class CragPage:
|
|||
class CragQuestion:
|
||||
"""One row of CRAG (Tasks 1 & 2)."""
|
||||
|
||||
qid: str # synthesised "C00000".."C02705"
|
||||
qid: str # synthesised "C00000".."C02705"
|
||||
interaction_id: str
|
||||
query_time: str
|
||||
query: str
|
||||
|
|
@ -89,9 +89,9 @@ class CragQuestion:
|
|||
domain: str
|
||||
question_type: str
|
||||
static_or_dynamic: str
|
||||
popularity: str # may be "" for web-sourced questions
|
||||
split: int # 0=validation, 1=public_test
|
||||
raw_index: int # row index in the source JSONL
|
||||
popularity: str # may be "" for web-sourced questions
|
||||
split: int # 0=validation, 1=public_test
|
||||
raw_index: int # row index in the source JSONL
|
||||
pages: list[CragPage] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
|
|
@ -166,16 +166,19 @@ def _parse_pages(raw_search_results: Any) -> list[CragPage]:
|
|||
if not url or not html.strip():
|
||||
# No URL or empty HTML => useless for retrieval.
|
||||
continue
|
||||
pages.append(CragPage(
|
||||
page_name=str(entry.get("page_name") or "").strip(),
|
||||
page_url=url,
|
||||
page_snippet=str(entry.get("page_snippet") or "").strip(),
|
||||
page_html=html,
|
||||
page_last_modified=(
|
||||
str(entry.get("page_last_modified")).strip()
|
||||
if entry.get("page_last_modified") else None
|
||||
),
|
||||
))
|
||||
pages.append(
|
||||
CragPage(
|
||||
page_name=str(entry.get("page_name") or "").strip(),
|
||||
page_url=url,
|
||||
page_snippet=str(entry.get("page_snippet") or "").strip(),
|
||||
page_html=html,
|
||||
page_last_modified=(
|
||||
str(entry.get("page_last_modified")).strip()
|
||||
if entry.get("page_last_modified")
|
||||
else None
|
||||
),
|
||||
)
|
||||
)
|
||||
return pages
|
||||
|
||||
|
||||
|
|
@ -217,21 +220,23 @@ def iter_questions(jsonl_bz2_path: Path) -> list[CragQuestion]:
|
|||
continue
|
||||
interaction_id = str(row.get("interaction_id") or "").strip()
|
||||
pages = _parse_pages(row.get("search_results"))
|
||||
out.append(CragQuestion(
|
||||
qid=f"C{raw_idx:05d}",
|
||||
interaction_id=interaction_id,
|
||||
query_time=str(row.get("query_time") or "").strip(),
|
||||
query=query,
|
||||
gold_answer=answer,
|
||||
alt_answers=_parse_alt_answers(row.get("alt_ans")),
|
||||
domain=str(row.get("domain") or "").strip().lower(),
|
||||
question_type=str(row.get("question_type") or "").strip().lower(),
|
||||
static_or_dynamic=str(row.get("static_or_dynamic") or "").strip().lower(),
|
||||
popularity=str(row.get("popularity") or "").strip().lower(),
|
||||
split=int(row.get("split") or 0),
|
||||
raw_index=raw_idx,
|
||||
pages=pages,
|
||||
))
|
||||
out.append(
|
||||
CragQuestion(
|
||||
qid=f"C{raw_idx:05d}",
|
||||
interaction_id=interaction_id,
|
||||
query_time=str(row.get("query_time") or "").strip(),
|
||||
query=query,
|
||||
gold_answer=answer,
|
||||
alt_answers=_parse_alt_answers(row.get("alt_ans")),
|
||||
domain=str(row.get("domain") or "").strip().lower(),
|
||||
question_type=str(row.get("question_type") or "").strip().lower(),
|
||||
static_or_dynamic=str(row.get("static_or_dynamic") or "").strip().lower(),
|
||||
popularity=str(row.get("popularity") or "").strip().lower(),
|
||||
split=int(row.get("split") or 0),
|
||||
raw_index=raw_idx,
|
||||
pages=pages,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -58,10 +58,10 @@ class CragGradeResult:
|
|||
"""One graded (pred, gold) pair under CRAG's 3-class rubric."""
|
||||
|
||||
grade: GradeClass
|
||||
score: int # +1 / 0 / -1
|
||||
method: str # exact, numeric, substring, refusal,
|
||||
# false_premise_correct, false_premise_miss,
|
||||
# llm_judge, lexical_miss, ...
|
||||
score: int # +1 / 0 / -1
|
||||
method: str # exact, numeric, substring, refusal,
|
||||
# false_premise_correct, false_premise_miss,
|
||||
# llm_judge, lexical_miss, ...
|
||||
normalised_pred: str = ""
|
||||
normalised_gold: str = ""
|
||||
judge_rationale: str = ""
|
||||
|
|
@ -112,10 +112,27 @@ def _normalise(s: str) -> str:
|
|||
|
||||
|
||||
_WORD_NUMBERS = {
|
||||
"zero": 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5,
|
||||
"six": 6, "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11,
|
||||
"twelve": 12, "thirteen": 13, "fourteen": 14, "fifteen": 15, "sixteen": 16,
|
||||
"seventeen": 17, "eighteen": 18, "nineteen": 19, "twenty": 20,
|
||||
"zero": 0,
|
||||
"one": 1,
|
||||
"two": 2,
|
||||
"three": 3,
|
||||
"four": 4,
|
||||
"five": 5,
|
||||
"six": 6,
|
||||
"seven": 7,
|
||||
"eight": 8,
|
||||
"nine": 9,
|
||||
"ten": 10,
|
||||
"eleven": 11,
|
||||
"twelve": 12,
|
||||
"thirteen": 13,
|
||||
"fourteen": 14,
|
||||
"fifteen": 15,
|
||||
"sixteen": 16,
|
||||
"seventeen": 17,
|
||||
"eighteen": 18,
|
||||
"nineteen": 19,
|
||||
"twenty": 20,
|
||||
}
|
||||
|
||||
_NUMERIC_RE = re.compile(r"-?\d+(?:[.,]\d+)?")
|
||||
|
|
@ -274,8 +291,11 @@ def grade_deterministic(
|
|||
continue
|
||||
if n_pred == cand_norm:
|
||||
return CragGradeResult(
|
||||
grade="correct", score=1, method="exact",
|
||||
normalised_pred=n_pred, normalised_gold=cand_norm,
|
||||
grade="correct",
|
||||
score=1,
|
||||
method="exact",
|
||||
normalised_pred=n_pred,
|
||||
normalised_gold=cand_norm,
|
||||
)
|
||||
p_num = _maybe_number(pred)
|
||||
c_num = _maybe_number(candidate)
|
||||
|
|
@ -289,21 +309,30 @@ def grade_deterministic(
|
|||
tol = abs(c_num) * 0.01
|
||||
if abs(p_num - c_num) <= tol:
|
||||
return CragGradeResult(
|
||||
grade="correct", score=1, method="numeric",
|
||||
normalised_pred=n_pred, normalised_gold=cand_norm,
|
||||
grade="correct",
|
||||
score=1,
|
||||
method="numeric",
|
||||
normalised_pred=n_pred,
|
||||
normalised_gold=cand_norm,
|
||||
)
|
||||
# Numeric question with different numbers — keep looking
|
||||
# at other candidates rather than declaring miss now;
|
||||
# alt answers may include word forms that pass.
|
||||
if _whole_word_substring(n_pred, cand_norm):
|
||||
return CragGradeResult(
|
||||
grade="correct", score=1, method="substring",
|
||||
normalised_pred=n_pred, normalised_gold=cand_norm,
|
||||
grade="correct",
|
||||
score=1,
|
||||
method="substring",
|
||||
normalised_pred=n_pred,
|
||||
normalised_gold=cand_norm,
|
||||
)
|
||||
if _whole_word_substring(cand_norm, n_pred) and len(n_pred) >= 3:
|
||||
return CragGradeResult(
|
||||
grade="correct", score=1, method="substring_reverse",
|
||||
normalised_pred=n_pred, normalised_gold=cand_norm,
|
||||
grade="correct",
|
||||
score=1,
|
||||
method="substring_reverse",
|
||||
normalised_pred=n_pred,
|
||||
normalised_gold=cand_norm,
|
||||
)
|
||||
|
||||
return CragGradeResult(
|
||||
|
|
@ -326,21 +355,21 @@ _JUDGE_SYSTEM = (
|
|||
"answer (and any alternative valid answers), and a model's "
|
||||
"prediction, classify the prediction into exactly one of three "
|
||||
"categories:\n\n"
|
||||
"* \"correct\" — the prediction expresses the same factual "
|
||||
'* "correct" — the prediction expresses the same factual '
|
||||
"content as the gold answer (paraphrasing OK; numbers as words "
|
||||
"OK; partial-but-correct names OK; non-contradictory extra "
|
||||
"detail OK).\n"
|
||||
"* \"missing\" — the prediction explicitly refuses, says \"I "
|
||||
'* "missing" — the prediction explicitly refuses, says "I '
|
||||
"don't know\", says there is insufficient information, or hedges "
|
||||
"without committing.\n"
|
||||
"* \"incorrect\" — the prediction commits to a fact that is "
|
||||
'* "incorrect" — the prediction commits to a fact that is '
|
||||
"different from the gold answer, or fails to flag a false "
|
||||
"premise when the question contains one.\n\n"
|
||||
"Special case: if the question contains a false premise and the "
|
||||
"gold answer says so, then a prediction that flags the false "
|
||||
"premise is \"correct\".\n\n"
|
||||
'premise is "correct".\n\n'
|
||||
"Respond with ONLY a JSON object on a single line:\n"
|
||||
'{\"grade\": \"correct\"|\"missing\"|\"incorrect\", \"rationale\": \"<one short sentence>\"}'
|
||||
'{"grade": "correct"|"missing"|"incorrect", "rationale": "<one short sentence>"}'
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -444,15 +473,17 @@ def _parse_judge_response(text: str) -> tuple[GradeClass, str]:
|
|||
|
||||
# Methods that should *not* trigger the LLM judge — the deterministic
|
||||
# verdict is conclusive (refusal, exact match, numeric mismatch, etc.).
|
||||
_TERMINAL_METHODS = frozenset({
|
||||
"refusal",
|
||||
"exact",
|
||||
"numeric",
|
||||
"substring",
|
||||
"substring_reverse",
|
||||
"false_premise_flagged",
|
||||
"empty_gold",
|
||||
})
|
||||
_TERMINAL_METHODS = frozenset(
|
||||
{
|
||||
"refusal",
|
||||
"exact",
|
||||
"numeric",
|
||||
"substring",
|
||||
"substring_reverse",
|
||||
"false_premise_flagged",
|
||||
"empty_gold",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def grade_with_judge(
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ class ExtractionResult:
|
|||
"""Outcome of converting one HTML blob to plain markdown."""
|
||||
|
||||
text: str
|
||||
method: str # "trafilatura" | "fallback_strip" | "empty"
|
||||
method: str # "trafilatura" | "fallback_strip" | "empty"
|
||||
n_chars: int
|
||||
|
||||
@property
|
||||
|
|
@ -94,11 +94,30 @@ class _StripHTMLParser(HTMLParser):
|
|||
"""
|
||||
|
||||
_SKIP_TAGS = frozenset({"script", "style", "nav", "header", "footer", "aside", "svg"})
|
||||
_BLOCK_TAGS = frozenset({
|
||||
"p", "div", "section", "article", "li", "ul", "ol",
|
||||
"h1", "h2", "h3", "h4", "h5", "h6", "br", "tr",
|
||||
"td", "th", "table", "blockquote", "pre",
|
||||
})
|
||||
_BLOCK_TAGS = frozenset(
|
||||
{
|
||||
"p",
|
||||
"div",
|
||||
"section",
|
||||
"article",
|
||||
"li",
|
||||
"ul",
|
||||
"ol",
|
||||
"h1",
|
||||
"h2",
|
||||
"h3",
|
||||
"h4",
|
||||
"h5",
|
||||
"h6",
|
||||
"br",
|
||||
"tr",
|
||||
"td",
|
||||
"th",
|
||||
"table",
|
||||
"blockquote",
|
||||
"pre",
|
||||
}
|
||||
)
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(convert_charrefs=True)
|
||||
|
|
|
|||
|
|
@ -158,7 +158,10 @@ def _materialise_pages(
|
|||
|
||||
logger.info(
|
||||
"CRAG page extraction: %s; empty=%d, total_files=%d across %d questions",
|
||||
method_counts, n_empty, len(file_to_url), len(qid_to_files),
|
||||
method_counts,
|
||||
n_empty,
|
||||
len(file_to_url),
|
||||
len(qid_to_files),
|
||||
)
|
||||
return qid_to_files, file_to_url
|
||||
|
||||
|
|
@ -215,8 +218,10 @@ async def _upload_pages(
|
|||
name_to_id[f"{s.title}.md"] = s.document_id
|
||||
logger.info(
|
||||
"CRAG upload batch %d-%d: %d new, %d duplicate",
|
||||
batch_start, batch_start + len(batch),
|
||||
len(result.document_ids), len(result.duplicate_document_ids),
|
||||
batch_start,
|
||||
batch_start + len(batch),
|
||||
len(result.document_ids),
|
||||
len(result.duplicate_document_ids),
|
||||
)
|
||||
return name_to_id
|
||||
|
||||
|
|
@ -243,24 +248,26 @@ def _resolve_question_doc_ids(
|
|||
doc_ids.append(doc_id)
|
||||
else:
|
||||
missing.append(fn)
|
||||
rows.append({
|
||||
"qid": q.qid,
|
||||
"interaction_id": q.interaction_id,
|
||||
"raw_index": q.raw_index,
|
||||
"question": q.query,
|
||||
"gold_answer": q.gold_answer,
|
||||
"alt_answers": list(q.alt_answers),
|
||||
"domain": q.domain,
|
||||
"question_type": q.question_type,
|
||||
"static_or_dynamic": q.static_or_dynamic,
|
||||
"popularity": q.popularity,
|
||||
"query_time": q.query_time,
|
||||
"split": q.split,
|
||||
"page_filenames": filenames,
|
||||
"document_ids": doc_ids,
|
||||
"missing_pages": missing,
|
||||
"n_pages": len(filenames),
|
||||
})
|
||||
rows.append(
|
||||
{
|
||||
"qid": q.qid,
|
||||
"interaction_id": q.interaction_id,
|
||||
"raw_index": q.raw_index,
|
||||
"question": q.query,
|
||||
"gold_answer": q.gold_answer,
|
||||
"alt_answers": list(q.alt_answers),
|
||||
"domain": q.domain,
|
||||
"question_type": q.question_type,
|
||||
"static_or_dynamic": q.static_or_dynamic,
|
||||
"popularity": q.popularity,
|
||||
"query_time": q.query_time,
|
||||
"split": q.split,
|
||||
"page_filenames": filenames,
|
||||
"document_ids": doc_ids,
|
||||
"missing_pages": missing,
|
||||
"n_pages": len(filenames),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
|
|
@ -305,7 +312,7 @@ async def run_ingest(
|
|||
settings = settings or IngestSettings(
|
||||
use_vision_llm=False,
|
||||
processing_mode="basic",
|
||||
)
|
||||
)
|
||||
bench_dir = ctx.benchmark_data_dir()
|
||||
pages_dir = bench_dir / "pages"
|
||||
raw_cache = bench_dir / ".raw_cache"
|
||||
|
|
@ -336,10 +343,13 @@ async def run_ingest(
|
|||
n_pages_total = sum(len(q.pages) for q in questions)
|
||||
logger.info(
|
||||
"CRAG: extracting up to %d pages across %d questions ...",
|
||||
n_pages_total, len(questions),
|
||||
n_pages_total,
|
||||
len(questions),
|
||||
)
|
||||
qid_to_files, file_to_url = _materialise_pages(
|
||||
questions, pages_dir=pages_dir, overwrite=overwrite_extract,
|
||||
questions,
|
||||
pages_dir=pages_dir,
|
||||
overwrite=overwrite_extract,
|
||||
)
|
||||
n_pages_extracted = sum(len(v) for v in qid_to_files.values())
|
||||
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ _BASE_INSTRUCTIONS = (
|
|||
"is factually wrong), say so explicitly in your final answer "
|
||||
"rather than answering as if the premise were true.\n"
|
||||
"2. If you are not confident in an answer, prefer saying \"I don't "
|
||||
"know\" over guessing. A wrong commit is penalised more than a "
|
||||
'know" over guessing. A wrong commit is penalised more than a '
|
||||
"refusal.\n"
|
||||
"3. Keep the final answer short — a name, a number, a date, a "
|
||||
"phrase. Do not repeat the question.\n\n"
|
||||
|
|
@ -125,9 +125,7 @@ def build_long_context_prompt(
|
|||
if len(body) > per_page_char_cap:
|
||||
body = body[:per_page_char_cap].rstrip() + "\n[...truncated...]"
|
||||
title_clean = (title or f"page_{idx}").strip().replace("\n", " ")
|
||||
blocks.append(
|
||||
f"--- PAGE {idx}: {title_clean} ---\n{body}\n"
|
||||
)
|
||||
blocks.append(f"--- PAGE {idx}: {title_clean} ---\n{body}\n")
|
||||
contexts_block = "\n".join(blocks) if blocks else "(no pages retrieved)"
|
||||
return _LONG_CONTEXT_TEMPLATE.format(
|
||||
instructions=_BASE_INSTRUCTIONS,
|
||||
|
|
|
|||
|
|
@ -125,21 +125,23 @@ def _filter_questions(
|
|||
continue
|
||||
if qtype_filter and qtype_filter not in qtype:
|
||||
continue
|
||||
out.append(CragRunnerQuestion(
|
||||
qid=str(row.get("qid") or "").strip(),
|
||||
raw_index=int(row.get("raw_index") or 0),
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=str(row.get("gold_answer") or "").strip(),
|
||||
alt_answers=list(row.get("alt_answers") or []),
|
||||
domain=domain,
|
||||
question_type=qtype,
|
||||
static_or_dynamic=str(row.get("static_or_dynamic") or "").lower(),
|
||||
popularity=str(row.get("popularity") or "").lower(),
|
||||
query_time=str(row.get("query_time") or "").strip(),
|
||||
page_filenames=list(row.get("page_filenames") or []),
|
||||
document_ids=list(row.get("document_ids") or []),
|
||||
missing_pages=list(row.get("missing_pages") or []),
|
||||
))
|
||||
out.append(
|
||||
CragRunnerQuestion(
|
||||
qid=str(row.get("qid") or "").strip(),
|
||||
raw_index=int(row.get("raw_index") or 0),
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=str(row.get("gold_answer") or "").strip(),
|
||||
alt_answers=list(row.get("alt_answers") or []),
|
||||
domain=domain,
|
||||
question_type=qtype,
|
||||
static_or_dynamic=str(row.get("static_or_dynamic") or "").lower(),
|
||||
popularity=str(row.get("popularity") or "").lower(),
|
||||
query_time=str(row.get("query_time") or "").strip(),
|
||||
page_filenames=list(row.get("page_filenames") or []),
|
||||
document_ids=list(row.get("document_ids") or []),
|
||||
missing_pages=list(row.get("missing_pages") or []),
|
||||
)
|
||||
)
|
||||
out.sort(key=lambda q: q.raw_index)
|
||||
if sample_n is not None and sample_n > 0:
|
||||
out = out[:sample_n]
|
||||
|
|
@ -190,15 +192,22 @@ class CragBenchmark:
|
|||
|
||||
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--n", dest="sample_n", type=int, default=None,
|
||||
"--n",
|
||||
dest="sample_n",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Run only the first N questions after filters.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--domain", dest="domain_filter", default=None,
|
||||
"--domain",
|
||||
dest="domain_filter",
|
||||
default=None,
|
||||
help="Filter to a single CRAG domain (finance|music|movie|sports|open).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--qtype", dest="qtype_filter", default=None,
|
||||
"--qtype",
|
||||
dest="qtype_filter",
|
||||
default=None,
|
||||
help=(
|
||||
"Filter to questions whose question_type contains this "
|
||||
"substring (case-insensitive). Examples: 'multi-hop', "
|
||||
|
|
@ -206,31 +215,46 @@ class CragBenchmark:
|
|||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, default=4,
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Parallel question workers per arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-output-tokens", type=int, default=512,
|
||||
"--max-output-tokens",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Cap on completion length for the chat-completion arms.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--per-page-char-cap", dest="per_page_char_cap", type=int, default=12_000,
|
||||
"--per-page-char-cap",
|
||||
dest="per_page_char_cap",
|
||||
type=int,
|
||||
default=12_000,
|
||||
help="Long-context arm: max chars per page before truncation (default 12k).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-bare", dest="skip_bare", action="store_true",
|
||||
"--skip-bare",
|
||||
dest="skip_bare",
|
||||
action="store_true",
|
||||
help="Skip the bare-LLM arm (saves cost on re-runs).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-long-context", dest="skip_long_context", action="store_true",
|
||||
"--skip-long-context",
|
||||
dest="skip_long_context",
|
||||
action="store_true",
|
||||
help="Skip the long-context arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-surfsense", dest="skip_surfsense", action="store_true",
|
||||
"--skip-surfsense",
|
||||
dest="skip_surfsense",
|
||||
action="store_true",
|
||||
help="Skip the SurfSense arm (useful when iterating on the LLM arms only).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-mention-scope", dest="no_mention_scope", action="store_true",
|
||||
"--no-mention-scope",
|
||||
dest="no_mention_scope",
|
||||
action="store_true",
|
||||
help=(
|
||||
"SurfSense arm: don't pass mentioned_document_ids; let "
|
||||
"the agent retrieve over the entire SearchSpace. Default "
|
||||
|
|
@ -239,37 +263,56 @@ class CragBenchmark:
|
|||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-judge", dest="no_judge", action="store_true",
|
||||
"--no-judge",
|
||||
dest="no_judge",
|
||||
action="store_true",
|
||||
help="Disable the LLM-as-judge fallback grader.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--judge-model", dest="judge_model",
|
||||
"--judge-model",
|
||||
dest="judge_model",
|
||||
default="anthropic/claude-sonnet-4.5",
|
||||
help="OpenRouter slug for the LLM judge.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--judge-concurrency", dest="judge_concurrency", type=int, default=4,
|
||||
"--judge-concurrency",
|
||||
dest="judge_concurrency",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Parallel judge calls.",
|
||||
)
|
||||
# Ingest knobs
|
||||
parser.add_argument(
|
||||
"--n-questions", dest="n_questions", type=int, default=None,
|
||||
"--n-questions",
|
||||
dest="n_questions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="(ingest only) cap on number of questions to materialise + ingest.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload-batch-size", dest="upload_batch_size", type=int, default=16,
|
||||
"--upload-batch-size",
|
||||
dest="upload_batch_size",
|
||||
type=int,
|
||||
default=16,
|
||||
help="(ingest only) markdown files per fileupload call.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-upload", dest="skip_upload", action="store_true",
|
||||
"--skip-upload",
|
||||
dest="skip_upload",
|
||||
action="store_true",
|
||||
help="(ingest only) extract pages locally but don't push to SurfSense.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite-extract", dest="overwrite_extract", action="store_true",
|
||||
"--overwrite-extract",
|
||||
dest="overwrite_extract",
|
||||
action="store_true",
|
||||
help="(ingest only) re-run trafilatura even when cached markdown exists.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-seed", dest="sample_seed", type=int, default=17,
|
||||
"--sample-seed",
|
||||
dest="sample_seed",
|
||||
type=int,
|
||||
default=17,
|
||||
help="(ingest only) RNG seed for the stratified sample.",
|
||||
)
|
||||
add_ingest_settings_args(parser, defaults=_DEFAULT_INGEST_SETTINGS)
|
||||
|
|
@ -362,12 +405,14 @@ class CragBenchmark:
|
|||
if not api_key:
|
||||
logger.warning("CRAG: --no-judge implied (no OPENROUTER_API_KEY for judge)")
|
||||
else:
|
||||
judge = CragLlmJudge(config=CragJudgeConfig(
|
||||
api_key=api_key,
|
||||
model=judge_model,
|
||||
base_url=ctx.config.openrouter_base_url,
|
||||
concurrency=judge_concurrency,
|
||||
))
|
||||
judge = CragLlmJudge(
|
||||
config=CragJudgeConfig(
|
||||
api_key=api_key,
|
||||
model=judge_model,
|
||||
base_url=ctx.config.openrouter_base_url,
|
||||
concurrency=judge_concurrency,
|
||||
)
|
||||
)
|
||||
|
||||
run_timestamp = utc_iso_timestamp()
|
||||
run_dir = ctx.runs_dir(run_timestamp=run_timestamp)
|
||||
|
|
@ -393,29 +438,53 @@ class CragBenchmark:
|
|||
# internally concurrency-bounded.
|
||||
tasks: list[Any] = []
|
||||
if bare_arm is not None:
|
||||
tasks.append(_gather_with_limit((_bare_one(q) for q in questions), concurrency=concurrency))
|
||||
tasks.append(
|
||||
_gather_with_limit((_bare_one(q) for q in questions), concurrency=concurrency)
|
||||
)
|
||||
else:
|
||||
tasks.append(_make_skipped_results(questions, "bare_llm"))
|
||||
if long_context_arm is not None:
|
||||
tasks.append(_gather_with_limit((_long_context_one(q) for q in questions), concurrency=concurrency))
|
||||
tasks.append(
|
||||
_gather_with_limit(
|
||||
(_long_context_one(q) for q in questions), concurrency=concurrency
|
||||
)
|
||||
)
|
||||
else:
|
||||
tasks.append(_make_skipped_results(questions, "long_context"))
|
||||
if surf_arm is not None:
|
||||
tasks.append(_gather_with_limit((_surf_one(q) for q in questions), concurrency=concurrency))
|
||||
tasks.append(
|
||||
_gather_with_limit((_surf_one(q) for q in questions), concurrency=concurrency)
|
||||
)
|
||||
else:
|
||||
tasks.append(_make_skipped_results(questions, "surfsense"))
|
||||
|
||||
bare_results, long_context_results, surf_results = await asyncio.gather(*tasks)
|
||||
|
||||
bare_grades = await _grade_results(questions, bare_results, judge=judge) if bare_arm else _empty_grades(questions)
|
||||
lc_grades = await _grade_results(questions, long_context_results, judge=judge) if long_context_arm else _empty_grades(questions)
|
||||
surf_grades = await _grade_results(questions, surf_results, judge=judge) if surf_arm else _empty_grades(questions)
|
||||
bare_grades = (
|
||||
await _grade_results(questions, bare_results, judge=judge)
|
||||
if bare_arm
|
||||
else _empty_grades(questions)
|
||||
)
|
||||
lc_grades = (
|
||||
await _grade_results(questions, long_context_results, judge=judge)
|
||||
if long_context_arm
|
||||
else _empty_grades(questions)
|
||||
)
|
||||
surf_grades = (
|
||||
await _grade_results(questions, surf_results, judge=judge)
|
||||
if surf_arm
|
||||
else _empty_grades(questions)
|
||||
)
|
||||
|
||||
with raw_path.open("w", encoding="utf-8") as fh:
|
||||
for q, b_res, l_res, s_res, b_g, l_g, s_g in zip(
|
||||
questions,
|
||||
bare_results, long_context_results, surf_results,
|
||||
bare_grades, lc_grades, surf_grades,
|
||||
bare_results,
|
||||
long_context_results,
|
||||
surf_results,
|
||||
bare_grades,
|
||||
lc_grades,
|
||||
surf_grades,
|
||||
strict=False,
|
||||
):
|
||||
meta = {
|
||||
|
|
@ -431,18 +500,29 @@ class CragBenchmark:
|
|||
"alt_answers": q.alt_answers,
|
||||
}
|
||||
for res, grade in (
|
||||
(b_res, b_g), (l_res, l_g), (s_res, s_g),
|
||||
(b_res, b_g),
|
||||
(l_res, l_g),
|
||||
(s_res, s_g),
|
||||
):
|
||||
fh.write(json.dumps({
|
||||
**meta,
|
||||
**res.to_jsonl(),
|
||||
"graded": grade.to_dict(),
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
**meta,
|
||||
**res.to_jsonl(),
|
||||
"graded": grade.to_dict(),
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
metrics = _compute_metrics(
|
||||
questions=questions,
|
||||
bare_results=bare_results, long_context_results=long_context_results, surf_results=surf_results,
|
||||
bare_grades=bare_grades, lc_grades=lc_grades, surf_grades=surf_grades,
|
||||
bare_results=bare_results,
|
||||
long_context_results=long_context_results,
|
||||
surf_results=surf_results,
|
||||
bare_grades=bare_grades,
|
||||
lc_grades=lc_grades,
|
||||
surf_grades=surf_grades,
|
||||
arms_active={
|
||||
"bare_llm": bare_arm is not None,
|
||||
"long_context": long_context_arm is not None,
|
||||
|
|
@ -481,13 +561,18 @@ class CragBenchmark:
|
|||
|
||||
manifest_path = run_dir / "run_artifact.json"
|
||||
manifest_path.write_text(
|
||||
json.dumps({
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
}, indent=2, sort_keys=True) + "\n",
|
||||
json.dumps(
|
||||
{
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
},
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return artifact
|
||||
|
|
@ -547,7 +632,9 @@ class CragBenchmark:
|
|||
|
||||
body_lines.append("- Headline truthfulness scores (CRAG paper rubric):")
|
||||
for label, key in (
|
||||
("Bare LLM", "bare_llm"), ("Long-Context", "long_context"), ("SurfSense", "surfsense"),
|
||||
("Bare LLM", "bare_llm"),
|
||||
("Long-Context", "long_context"),
|
||||
("SurfSense", "surfsense"),
|
||||
):
|
||||
d = m.get(key, {})
|
||||
body_lines.append(
|
||||
|
|
@ -583,9 +670,7 @@ class CragBenchmark:
|
|||
for arm in ("bare_llm", "long_context", "surfsense"):
|
||||
if arm not in row:
|
||||
continue
|
||||
pieces.append(
|
||||
f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}"
|
||||
)
|
||||
pieces.append(f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}")
|
||||
body_lines.append(" ".join(pieces))
|
||||
|
||||
if per_qtype:
|
||||
|
|
@ -596,9 +681,7 @@ class CragBenchmark:
|
|||
for arm in ("bare_llm", "long_context", "surfsense"):
|
||||
if arm not in row:
|
||||
continue
|
||||
pieces.append(
|
||||
f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}"
|
||||
)
|
||||
pieces.append(f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}")
|
||||
body_lines.append(" ".join(pieces))
|
||||
|
||||
return ReportSection(
|
||||
|
|
@ -669,32 +752,31 @@ async def _grade_results(
|
|||
rows: list[CragGradeRow] = []
|
||||
for q, r in zip(questions, results, strict=False):
|
||||
pred = extract_freeform_answer(r.raw_text or "")
|
||||
rows.append(CragGradeRow(
|
||||
qid=q.qid,
|
||||
question=q.question,
|
||||
gold=q.gold_answer,
|
||||
alt_answers=q.alt_answers,
|
||||
pred=pred,
|
||||
question_type=q.question_type,
|
||||
))
|
||||
rows.append(
|
||||
CragGradeRow(
|
||||
qid=q.qid,
|
||||
question=q.question,
|
||||
gold=q.gold_answer,
|
||||
alt_answers=q.alt_answers,
|
||||
pred=pred,
|
||||
question_type=q.question_type,
|
||||
)
|
||||
)
|
||||
return await grade_many(rows=rows, judge=judge)
|
||||
|
||||
|
||||
def _empty_grades(questions: list[CragRunnerQuestion]) -> list[CragGradeResult]:
|
||||
return [
|
||||
CragGradeResult(grade="missing", score=0, method="skipped_arm")
|
||||
for _ in questions
|
||||
]
|
||||
return [CragGradeResult(grade="missing", score=0, method="skipped_arm") for _ in questions]
|
||||
|
||||
|
||||
async def _make_skipped_results(
|
||||
questions: list[CragRunnerQuestion], arm_name: str,
|
||||
questions: list[CragRunnerQuestion],
|
||||
arm_name: str,
|
||||
) -> list[ArmResult]:
|
||||
"""Stand-in results so downstream code can assume parallel lists."""
|
||||
|
||||
return [
|
||||
ArmResult(arm=arm_name, question_id=q.qid, raw_text="", error="skipped")
|
||||
for q in questions
|
||||
ArmResult(arm=arm_name, question_id=q.qid, raw_text="", error="skipped") for q in questions
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -776,20 +858,41 @@ def _compute_metrics(
|
|||
|
||||
deltas: dict[str, Any] = {}
|
||||
for label, ref_correct, ref_t, chal_correct, chal_t, both_active in (
|
||||
("surfsense_vs_bare", bare_correct, bare_t, surf_correct, surf_t,
|
||||
arms_active.get("bare_llm") and arms_active.get("surfsense")),
|
||||
("surfsense_vs_long_context", lc_correct, lc_t, surf_correct, surf_t,
|
||||
arms_active.get("long_context") and arms_active.get("surfsense")),
|
||||
("long_context_vs_bare", bare_correct, bare_t, lc_correct, lc_t,
|
||||
arms_active.get("bare_llm") and arms_active.get("long_context")),
|
||||
(
|
||||
"surfsense_vs_bare",
|
||||
bare_correct,
|
||||
bare_t,
|
||||
surf_correct,
|
||||
surf_t,
|
||||
arms_active.get("bare_llm") and arms_active.get("surfsense"),
|
||||
),
|
||||
(
|
||||
"surfsense_vs_long_context",
|
||||
lc_correct,
|
||||
lc_t,
|
||||
surf_correct,
|
||||
surf_t,
|
||||
arms_active.get("long_context") and arms_active.get("surfsense"),
|
||||
),
|
||||
(
|
||||
"long_context_vs_bare",
|
||||
bare_correct,
|
||||
bare_t,
|
||||
lc_correct,
|
||||
lc_t,
|
||||
arms_active.get("bare_llm") and arms_active.get("long_context"),
|
||||
),
|
||||
):
|
||||
if not both_active:
|
||||
continue
|
||||
mc = mcnemar_test(ref_correct, chal_correct)
|
||||
boot = bootstrap_delta_ci(ref_correct, chal_correct, n_resamples=2000)
|
||||
deltas[label] = {
|
||||
"accuracy_pp": 100.0 * (sum(chal_correct) - sum(ref_correct)) / max(1, len(chal_correct)),
|
||||
"truthfulness_score_pp": 100.0 * (chal_t["truthfulness_score"] - ref_t["truthfulness_score"]),
|
||||
"accuracy_pp": 100.0
|
||||
* (sum(chal_correct) - sum(ref_correct))
|
||||
/ max(1, len(chal_correct)),
|
||||
"truthfulness_score_pp": 100.0
|
||||
* (chal_t["truthfulness_score"] - ref_t["truthfulness_score"]),
|
||||
"mcnemar_p_value": mc.p_value,
|
||||
"mcnemar_method": mc.method,
|
||||
"mcnemar_b_ref_only": mc.b,
|
||||
|
|
@ -800,12 +903,18 @@ def _compute_metrics(
|
|||
out["deltas"] = deltas
|
||||
|
||||
out["per_domain"] = _per_facet_truthfulness(
|
||||
questions, bare_grades, lc_grades, surf_grades,
|
||||
questions,
|
||||
bare_grades,
|
||||
lc_grades,
|
||||
surf_grades,
|
||||
arms_active=arms_active,
|
||||
key_fn=lambda q: q.domain or "(unspecified)",
|
||||
)
|
||||
out["per_question_type"] = _per_facet_truthfulness(
|
||||
questions, bare_grades, lc_grades, surf_grades,
|
||||
questions,
|
||||
bare_grades,
|
||||
lc_grades,
|
||||
surf_grades,
|
||||
arms_active=arms_active,
|
||||
key_fn=lambda q: q.question_type or "(unspecified)",
|
||||
)
|
||||
|
|
@ -867,11 +976,11 @@ def _arm_summary_lines(d: dict[str, Any], *, indent: str) -> str:
|
|||
high = d.get("ci_high", 0.0)
|
||||
lines = [
|
||||
f"{indent}- Accuracy: {acc * 100:.1f}% (Wilson 95% CI: {low * 100:.1f}% – {high * 100:.1f}%)",
|
||||
f"{indent}- 3-class: correct={d.get('correct_rate', 0)*100:.1f}%, "
|
||||
f"missing={d.get('missing_rate', 0)*100:.1f}%, "
|
||||
f"incorrect={d.get('incorrect_rate', 0)*100:.1f}%",
|
||||
f"{indent}- 3-class: correct={d.get('correct_rate', 0) * 100:.1f}%, "
|
||||
f"missing={d.get('missing_rate', 0) * 100:.1f}%, "
|
||||
f"incorrect={d.get('incorrect_rate', 0) * 100:.1f}%",
|
||||
f"{indent}- Truthfulness score (correct - incorrect)/total: "
|
||||
f"{d.get('truthfulness_score', 0)*100:+.1f}%",
|
||||
f"{d.get('truthfulness_score', 0) * 100:+.1f}%",
|
||||
f"{indent}- Cost / question: ${_dollars(d.get('cost_micros_mean'))} (mean), "
|
||||
f"${_dollars(d.get('cost_micros_median'))} (median)",
|
||||
f"{indent}- Latency: p50 {_ms_to_s(d.get('latency_ms_median'))}, "
|
||||
|
|
@ -916,7 +1025,7 @@ def _pct(value: Any) -> str:
|
|||
if value is None:
|
||||
return "?"
|
||||
try:
|
||||
return f"{float(value)*100:.1f}%"
|
||||
return f"{float(value) * 100:.1f}%"
|
||||
except (TypeError, ValueError):
|
||||
return "?"
|
||||
|
||||
|
|
@ -925,7 +1034,7 @@ def _signed_pct(value: Any) -> str:
|
|||
if value is None:
|
||||
return "?"
|
||||
try:
|
||||
return f"{float(value)*100:+.1f}%"
|
||||
return f"{float(value) * 100:+.1f}%"
|
||||
except (TypeError, ValueError):
|
||||
return "?"
|
||||
|
||||
|
|
|
|||
|
|
@ -51,12 +51,12 @@ def _hf_hub_download(*args: Any, **kwargs: Any) -> str:
|
|||
class FramesQuestion:
|
||||
"""One row of FRAMES (post-parse)."""
|
||||
|
||||
qid: str # synthesised "Q000" .. "Q823"
|
||||
qid: str # synthesised "Q000" .. "Q823"
|
||||
question: str
|
||||
gold_answer: str
|
||||
wiki_urls: list[str] # deduped, in original order
|
||||
reasoning_types: list[str] # split on "|"
|
||||
raw_index: int # row index from the TSV (for debugging)
|
||||
wiki_urls: list[str] # deduped, in original order
|
||||
reasoning_types: list[str] # split on "|"
|
||||
raw_index: int # row index from the TSV (for debugging)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
|
|
@ -146,14 +146,16 @@ def load_questions(tsv_path: Path) -> list[FramesQuestion]:
|
|||
if val and val not in urls:
|
||||
urls.append(val)
|
||||
reasoning = _parse_reasoning_types(row.get("reasoning_types"))
|
||||
out.append(FramesQuestion(
|
||||
qid=f"Q{int(raw_idx):03d}",
|
||||
question=prompt,
|
||||
gold_answer=answer,
|
||||
wiki_urls=urls,
|
||||
reasoning_types=reasoning,
|
||||
raw_index=int(raw_idx),
|
||||
))
|
||||
out.append(
|
||||
FramesQuestion(
|
||||
qid=f"Q{int(raw_idx):03d}",
|
||||
question=prompt,
|
||||
gold_answer=answer,
|
||||
wiki_urls=urls,
|
||||
reasoning_types=reasoning,
|
||||
raw_index=int(raw_idx),
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -90,10 +90,27 @@ def _normalise(s: str) -> str:
|
|||
|
||||
|
||||
_WORD_NUMBERS = {
|
||||
"zero": 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5,
|
||||
"six": 6, "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11,
|
||||
"twelve": 12, "thirteen": 13, "fourteen": 14, "fifteen": 15, "sixteen": 16,
|
||||
"seventeen": 17, "eighteen": 18, "nineteen": 19, "twenty": 20,
|
||||
"zero": 0,
|
||||
"one": 1,
|
||||
"two": 2,
|
||||
"three": 3,
|
||||
"four": 4,
|
||||
"five": 5,
|
||||
"six": 6,
|
||||
"seven": 7,
|
||||
"eight": 8,
|
||||
"nine": 9,
|
||||
"ten": 10,
|
||||
"eleven": 11,
|
||||
"twelve": 12,
|
||||
"thirteen": 13,
|
||||
"fourteen": 14,
|
||||
"fifteen": 15,
|
||||
"sixteen": 16,
|
||||
"seventeen": 17,
|
||||
"eighteen": 18,
|
||||
"nineteen": 19,
|
||||
"twenty": 20,
|
||||
}
|
||||
|
||||
_NUMERIC_RE = re.compile(r"-?\d+(?:[.,]\d+)?")
|
||||
|
|
@ -194,7 +211,7 @@ _JUDGE_SYSTEM = (
|
|||
"expresses a different fact, omits the central answer, or hedges "
|
||||
"without committing.\n\n"
|
||||
"Respond with ONLY a JSON object on a single line:\n"
|
||||
'{\"correct\": true|false, \"rationale\": \"<one short sentence>\"}'
|
||||
'{"correct": true|false, "rationale": "<one short sentence>"}'
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -324,10 +341,7 @@ async def grade_many(
|
|||
|
||||
if not rows:
|
||||
return []
|
||||
coros = [
|
||||
grade_with_judge(pred=p, gold=g, question=q, judge=judge)
|
||||
for _qid, q, g, p in rows
|
||||
]
|
||||
coros = [grade_with_judge(pred=p, gold=g, question=q, judge=judge) for _qid, q, g, p in rows]
|
||||
return list(await asyncio.gather(*coros))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -160,8 +160,10 @@ async def _upload_markdowns(
|
|||
name_to_id[s.title] = s.document_id
|
||||
logger.info(
|
||||
"FRAMES upload batch %d-%d: %d new, %d duplicate",
|
||||
batch_start, batch_start + len(batch),
|
||||
len(result.document_ids), len(result.duplicate_document_ids),
|
||||
batch_start,
|
||||
batch_start + len(batch),
|
||||
len(result.document_ids),
|
||||
len(result.duplicate_document_ids),
|
||||
)
|
||||
return name_to_id
|
||||
|
||||
|
|
@ -188,14 +190,16 @@ def _resolve_question_doc_ids(
|
|||
doc_id = name_to_id.get(stem) or name_to_id.get(article.markdown_path.name)
|
||||
if doc_id is not None and doc_id not in doc_ids:
|
||||
doc_ids.append(doc_id)
|
||||
rows.append({
|
||||
"qid": q.qid,
|
||||
"raw_index": q.raw_index,
|
||||
"n_wiki_urls": len(q.wiki_urls),
|
||||
"wiki_titles": titles,
|
||||
"document_ids": doc_ids,
|
||||
"missing_urls": missing,
|
||||
})
|
||||
rows.append(
|
||||
{
|
||||
"qid": q.qid,
|
||||
"raw_index": q.raw_index,
|
||||
"n_wiki_urls": len(q.wiki_urls),
|
||||
"wiki_titles": titles,
|
||||
"document_ids": doc_ids,
|
||||
"missing_urls": missing,
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
|
|
@ -238,7 +242,7 @@ async def run_ingest(
|
|||
settings = settings or IngestSettings(
|
||||
use_vision_llm=False,
|
||||
processing_mode="basic",
|
||||
)
|
||||
)
|
||||
bench_dir = ctx.benchmark_data_dir()
|
||||
wiki_cache = bench_dir / "wiki"
|
||||
wiki_cache.mkdir(parents=True, exist_ok=True)
|
||||
|
|
@ -250,8 +254,7 @@ async def run_ingest(
|
|||
questions = load_questions(tsv_path)
|
||||
if not questions:
|
||||
raise RuntimeError(
|
||||
"FRAMES test.tsv contained no parseable rows; upstream may "
|
||||
"have changed schema."
|
||||
"FRAMES test.tsv contained no parseable rows; upstream may have changed schema."
|
||||
)
|
||||
logger.info("FRAMES: parsed %d questions from %s", len(questions), tsv_path.name)
|
||||
if max_questions is not None and max_questions > 0:
|
||||
|
|
@ -269,19 +272,23 @@ async def run_ingest(
|
|||
unique_urls = list(seen_urls.keys())
|
||||
logger.info(
|
||||
"FRAMES: %d unique Wikipedia URLs across %d questions",
|
||||
len(unique_urls), len(questions),
|
||||
len(unique_urls),
|
||||
len(questions),
|
||||
)
|
||||
|
||||
# 3. Fetch (with cache).
|
||||
fetcher = WikiFetcher(cache_dir=wiki_cache, rate_limit_rps=fetch_rate_limit_rps)
|
||||
n_cached = sum(
|
||||
1 for url in unique_urls
|
||||
1
|
||||
for url in unique_urls
|
||||
if (wiki_cache / cache_filename_for_title(_safe_title(url))).exists()
|
||||
)
|
||||
fetched, missing_urls = await _fetch_articles(fetcher, unique_urls)
|
||||
logger.info(
|
||||
"FRAMES: fetched=%d, cache_hits=%d, missing=%d",
|
||||
len(fetched), n_cached, len(missing_urls),
|
||||
len(fetched),
|
||||
n_cached,
|
||||
len(missing_urls),
|
||||
)
|
||||
|
||||
# 4. Upload to SurfSense (deduped by filename).
|
||||
|
|
|
|||
|
|
@ -65,7 +65,7 @@ class FramesRunnerQuestion:
|
|||
question: str
|
||||
gold_answer: str
|
||||
reasoning_types: list[str]
|
||||
document_ids: list[int] # subset of corpus relevant to this Q (may be empty)
|
||||
document_ids: list[int] # subset of corpus relevant to this Q (may be empty)
|
||||
n_wiki_urls: int
|
||||
missing_urls: list[str]
|
||||
|
||||
|
|
@ -107,16 +107,18 @@ def _load_questions(
|
|||
reasoning = list(row.get("reasoning_types") or [])
|
||||
if reasoning_filter and reasoning_filter not in [r.lower() for r in reasoning]:
|
||||
continue
|
||||
out.append(FramesRunnerQuestion(
|
||||
qid=qid,
|
||||
raw_index=int(row.get("raw_index") or 0),
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=str(row.get("gold_answer") or "").strip(),
|
||||
reasoning_types=reasoning,
|
||||
document_ids=list(map_row.get("document_ids") or []),
|
||||
n_wiki_urls=int(map_row.get("n_wiki_urls") or 0),
|
||||
missing_urls=list(map_row.get("missing_urls") or []),
|
||||
))
|
||||
out.append(
|
||||
FramesRunnerQuestion(
|
||||
qid=qid,
|
||||
raw_index=int(row.get("raw_index") or 0),
|
||||
question=str(row.get("question") or "").strip(),
|
||||
gold_answer=str(row.get("gold_answer") or "").strip(),
|
||||
reasoning_types=reasoning,
|
||||
document_ids=list(map_row.get("document_ids") or []),
|
||||
n_wiki_urls=int(map_row.get("n_wiki_urls") or 0),
|
||||
missing_urls=list(map_row.get("missing_urls") or []),
|
||||
)
|
||||
)
|
||||
out.sort(key=lambda q: q.raw_index)
|
||||
if sample_n is not None and sample_n > 0:
|
||||
out = out[:sample_n]
|
||||
|
|
@ -166,7 +168,10 @@ class FramesBenchmark:
|
|||
|
||||
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--n", dest="sample_n", type=int, default=None,
|
||||
"--n",
|
||||
dest="sample_n",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Run only the first N questions after filters (default: all 824).",
|
||||
)
|
||||
parser.add_argument(
|
||||
|
|
@ -180,11 +185,15 @@ class FramesBenchmark:
|
|||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, default=4,
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Parallel question workers per arm.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--scope-mentions", dest="scope_mentions", action="store_true",
|
||||
"--scope-mentions",
|
||||
dest="scope_mentions",
|
||||
action="store_true",
|
||||
help=(
|
||||
"SurfSense arm: scope retrieval to the per-question "
|
||||
"document_ids (oracle-retrieval upper bound). Default "
|
||||
|
|
@ -192,11 +201,15 @@ class FramesBenchmark:
|
|||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-output-tokens", type=int, default=512,
|
||||
"--max-output-tokens",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Cap on completion length for both arms.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-judge", dest="no_judge", action="store_true",
|
||||
"--no-judge",
|
||||
dest="no_judge",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Disable LLM-as-judge fallback grading; use only the "
|
||||
"deterministic grader (faster but more pessimistic)."
|
||||
|
|
@ -217,19 +230,30 @@ class FramesBenchmark:
|
|||
)
|
||||
# Ingest-only knobs.
|
||||
parser.add_argument(
|
||||
"--max-questions", dest="max_questions", type=int, default=None,
|
||||
"--max-questions",
|
||||
dest="max_questions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="(ingest only) cap on number of questions to materialise + ingest.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload-batch-size", dest="upload_batch_size", type=int, default=16,
|
||||
"--upload-batch-size",
|
||||
dest="upload_batch_size",
|
||||
type=int,
|
||||
default=16,
|
||||
help="(ingest only) markdown files per fileupload call.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-upload", dest="skip_upload", action="store_true",
|
||||
"--skip-upload",
|
||||
dest="skip_upload",
|
||||
action="store_true",
|
||||
help="(ingest only) cache wiki articles locally but don't push to SurfSense.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fetch-rps", dest="fetch_rate_limit_rps", type=float, default=2.0,
|
||||
"--fetch-rps",
|
||||
dest="fetch_rate_limit_rps",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="(ingest only) max requests/second to the Wikipedia API.",
|
||||
)
|
||||
add_ingest_settings_args(parser, defaults=_DEFAULT_INGEST_SETTINGS)
|
||||
|
|
@ -270,21 +294,18 @@ class FramesBenchmark:
|
|||
|
||||
doc_map, ingest_settings = _load_doc_map(map_path)
|
||||
questions = _load_questions(
|
||||
questions_jsonl, doc_map,
|
||||
questions_jsonl,
|
||||
doc_map,
|
||||
sample_n=sample_n,
|
||||
reasoning_filter=reasoning_filter,
|
||||
)
|
||||
if not questions:
|
||||
raise RuntimeError(
|
||||
"No FRAMES questions matched the filters; broaden --reasoning/--n."
|
||||
)
|
||||
raise RuntimeError("No FRAMES questions matched the filters; broaden --reasoning/--n.")
|
||||
logger.info("FRAMES: scheduled %d questions", len(questions))
|
||||
|
||||
api_key = os.environ.get("OPENROUTER_API_KEY")
|
||||
if not api_key:
|
||||
raise RuntimeError(
|
||||
"OPENROUTER_API_KEY env var is required for the bare-LLM arm."
|
||||
)
|
||||
raise RuntimeError("OPENROUTER_API_KEY env var is required for the bare-LLM arm.")
|
||||
|
||||
bare_provider = OpenRouterChatProvider(
|
||||
api_key=api_key,
|
||||
|
|
@ -303,12 +324,14 @@ class FramesBenchmark:
|
|||
|
||||
judge: LlmJudge | None = None
|
||||
if not no_judge:
|
||||
judge = LlmJudge(config=JudgeConfig(
|
||||
api_key=api_key,
|
||||
model=judge_model,
|
||||
base_url=ctx.config.openrouter_base_url,
|
||||
concurrency=judge_concurrency,
|
||||
))
|
||||
judge = LlmJudge(
|
||||
config=JudgeConfig(
|
||||
api_key=api_key,
|
||||
model=judge_model,
|
||||
base_url=ctx.config.openrouter_base_url,
|
||||
concurrency=judge_concurrency,
|
||||
)
|
||||
)
|
||||
|
||||
run_timestamp = utc_iso_timestamp()
|
||||
run_dir = ctx.runs_dir(run_timestamp=run_timestamp)
|
||||
|
|
@ -318,9 +341,7 @@ class FramesBenchmark:
|
|||
return await bare_arm.answer(_make_bare_request(q, max_output_tokens))
|
||||
|
||||
async def _surf_one(q: FramesRunnerQuestion) -> ArmResult:
|
||||
return await surf_arm.answer(
|
||||
_make_surfsense_request(q, scope_mentions=scope_mentions)
|
||||
)
|
||||
return await surf_arm.answer(_make_surfsense_request(q, scope_mentions=scope_mentions))
|
||||
|
||||
bare_results, surf_results = await asyncio.gather(
|
||||
_gather_with_limit((_bare_one(q) for q in questions), concurrency=concurrency),
|
||||
|
|
@ -343,16 +364,26 @@ class FramesBenchmark:
|
|||
"n_missing_urls": len(q.missing_urls),
|
||||
"gold": q.gold_answer,
|
||||
}
|
||||
fh.write(json.dumps({
|
||||
**meta,
|
||||
**b_res.to_jsonl(),
|
||||
"graded": b_g.to_dict(),
|
||||
}) + "\n")
|
||||
fh.write(json.dumps({
|
||||
**meta,
|
||||
**s_res.to_jsonl(),
|
||||
"graded": s_g.to_dict(),
|
||||
}) + "\n")
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
**meta,
|
||||
**b_res.to_jsonl(),
|
||||
"graded": b_g.to_dict(),
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
fh.write(
|
||||
json.dumps(
|
||||
{
|
||||
**meta,
|
||||
**s_res.to_jsonl(),
|
||||
"graded": s_g.to_dict(),
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
metrics = _compute_metrics(questions, bare_results, surf_results, bare_grades, surf_grades)
|
||||
artifact = RunArtifact(
|
||||
|
|
@ -380,13 +411,18 @@ class FramesBenchmark:
|
|||
|
||||
manifest_path = run_dir / "run_artifact.json"
|
||||
manifest_path.write_text(
|
||||
json.dumps({
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
}, indent=2, sort_keys=True) + "\n",
|
||||
json.dumps(
|
||||
{
|
||||
"suite": self.suite,
|
||||
"benchmark": self.name,
|
||||
"raw_path": "raw.jsonl",
|
||||
"metrics": metrics,
|
||||
"extra": artifact.extra,
|
||||
},
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return artifact
|
||||
|
|
@ -451,8 +487,8 @@ class FramesBenchmark:
|
|||
for tag, vals in sorted(per_reasoning.items()):
|
||||
body_lines.append(
|
||||
f" - {tag}: SurfSense {_pp(vals.get('delta_accuracy_pp'))} pp "
|
||||
f"(n={vals.get('n')}, bare acc={vals.get('bare_accuracy', 0)*100:.1f}%, "
|
||||
f"surf acc={vals.get('surfsense_accuracy', 0)*100:.1f}%)"
|
||||
f"(n={vals.get('n')}, bare acc={vals.get('bare_accuracy', 0) * 100:.1f}%, "
|
||||
f"surf acc={vals.get('surfsense_accuracy', 0) * 100:.1f}%)"
|
||||
)
|
||||
|
||||
return ReportSection(
|
||||
|
|
@ -553,8 +589,7 @@ def _compute_metrics(
|
|||
"bare_accuracy": (sum(b_correct) / len(pairs)) if pairs else 0.0,
|
||||
"surfsense_accuracy": (sum(s_correct) / len(pairs)) if pairs else 0.0,
|
||||
"delta_accuracy_pp": (
|
||||
100.0 * (sum(s_correct) - sum(b_correct)) / len(pairs)
|
||||
if pairs else 0.0
|
||||
100.0 * (sum(s_correct) - sum(b_correct)) / len(pairs) if pairs else 0.0
|
||||
),
|
||||
}
|
||||
|
||||
|
|
@ -571,8 +606,12 @@ def _compute_metrics(
|
|||
"latency_ms_mean": bare_latency_agg.mean,
|
||||
"latency_ms_median": bare_latency_agg.median,
|
||||
"latency_ms_p95": bare_latency_agg.p95,
|
||||
"input_tokens_mean": (sum(bare_in_tokens) / len(bare_in_tokens)) if bare_in_tokens else 0.0,
|
||||
"output_tokens_mean": (sum(bare_out_tokens) / len(bare_out_tokens)) if bare_out_tokens else 0.0,
|
||||
"input_tokens_mean": (sum(bare_in_tokens) / len(bare_in_tokens))
|
||||
if bare_in_tokens
|
||||
else 0.0,
|
||||
"output_tokens_mean": (sum(bare_out_tokens) / len(bare_out_tokens))
|
||||
if bare_out_tokens
|
||||
else 0.0,
|
||||
},
|
||||
"surfsense": {
|
||||
**surf_acc.to_dict(),
|
||||
|
|
|
|||
|
|
@ -49,10 +49,10 @@ USER_AGENT = (
|
|||
class WikiArticle:
|
||||
"""One fetched article + metadata."""
|
||||
|
||||
title: str # canonical title returned by MW (post-redirect)
|
||||
source_url: str # the URL we were asked to fetch
|
||||
markdown_path: Path # where the cached body lives on disk
|
||||
n_chars: int # length of the body (post-prepend H1)
|
||||
title: str # canonical title returned by MW (post-redirect)
|
||||
source_url: str # the URL we were asked to fetch
|
||||
markdown_path: Path # where the cached body lives on disk
|
||||
n_chars: int # length of the body (post-prepend H1)
|
||||
redirected_from: str | None = None
|
||||
|
||||
|
||||
|
|
@ -168,10 +168,13 @@ class WikiFetcher:
|
|||
break
|
||||
except (httpx.HTTPError, RuntimeError) as exc:
|
||||
last_exc = exc
|
||||
wait = 1.0 * (2 ** attempt)
|
||||
wait = 1.0 * (2**attempt)
|
||||
logger.warning(
|
||||
"wiki fetch %r attempt %d failed: %s; retry in %.1fs",
|
||||
title, attempt + 1, exc, wait,
|
||||
title,
|
||||
attempt + 1,
|
||||
exc,
|
||||
wait,
|
||||
)
|
||||
await asyncio.sleep(wait)
|
||||
else:
|
||||
|
|
@ -217,10 +220,14 @@ class WikiFetcher:
|
|||
}
|
||||
headers = {"User-Agent": USER_AGENT, "Accept": "application/json"}
|
||||
if http is not None:
|
||||
response = await http.get(WIKI_API, params=params, headers=headers, timeout=self._timeout)
|
||||
response = await http.get(
|
||||
WIKI_API, params=params, headers=headers, timeout=self._timeout
|
||||
)
|
||||
else:
|
||||
async with httpx.AsyncClient(timeout=self._timeout) as client:
|
||||
response = await client.get(WIKI_API, params=params, headers=headers, timeout=self._timeout)
|
||||
response = await client.get(
|
||||
WIKI_API, params=params, headers=headers, timeout=self._timeout
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if "error" in data:
|
||||
|
|
|
|||
|
|
@ -52,9 +52,7 @@ async def test_acquire_token_local_mode_posts_desktop_login_json():
|
|||
200, json={"access_token": "T", "refresh_token": "R", "token_type": "bearer"}
|
||||
)
|
||||
)
|
||||
config = _make_config(
|
||||
surfsense_user_email="u@example.com", surfsense_user_password="pw"
|
||||
)
|
||||
config = _make_config(surfsense_user_email="u@example.com", surfsense_user_password="pw")
|
||||
bundle = await acquire_token(config)
|
||||
assert bundle.access_token == "T"
|
||||
assert bundle.refresh_token == "R"
|
||||
|
|
|
|||
|
|
@ -94,10 +94,18 @@ async def test_documents_status_parses_state(respx_mock, http):
|
|||
200,
|
||||
json={
|
||||
"items": [
|
||||
{"id": 1, "title": "a.pdf", "document_type": "FILE",
|
||||
"status": {"state": "ready", "reason": None}},
|
||||
{"id": 2, "title": "b.pdf", "document_type": "FILE",
|
||||
"status": {"state": "failed", "reason": "ETL boom"}},
|
||||
{
|
||||
"id": 1,
|
||||
"title": "a.pdf",
|
||||
"document_type": "FILE",
|
||||
"status": {"state": "ready", "reason": None},
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"title": "b.pdf",
|
||||
"document_type": "FILE",
|
||||
"status": {"state": "failed", "reason": "ETL boom"},
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
|
|
@ -137,14 +145,26 @@ async def test_documents_upload_returns_payload(respx_mock, http, tmp_path: Path
|
|||
async def test_documents_list_chunks_paginated(respx_mock, http):
|
||||
respx_mock.get("/api/v1/documents/5/chunks").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json={
|
||||
"items": [{"id": 1, "content": "a"}, {"id": 2, "content": "b"}],
|
||||
"total": 3, "page": 0, "page_size": 2, "has_more": True,
|
||||
}),
|
||||
httpx.Response(200, json={
|
||||
"items": [{"id": 3, "content": "c"}],
|
||||
"total": 3, "page": 1, "page_size": 2, "has_more": False,
|
||||
}),
|
||||
httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"items": [{"id": 1, "content": "a"}, {"id": 2, "content": "b"}],
|
||||
"total": 3,
|
||||
"page": 0,
|
||||
"page_size": 2,
|
||||
"has_more": True,
|
||||
},
|
||||
),
|
||||
httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"items": [{"id": 3, "content": "c"}],
|
||||
"total": 3,
|
||||
"page": 1,
|
||||
"page_size": 2,
|
||||
"has_more": False,
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
client = DocumentsClient(http, _BASE)
|
||||
|
|
@ -191,15 +211,17 @@ def _sse_body(events: list[dict]) -> bytes:
|
|||
@pytest.mark.asyncio
|
||||
@respx.mock(base_url=_BASE)
|
||||
async def test_ask_accumulates_text_deltas(respx_mock, http):
|
||||
body = _sse_body([
|
||||
{"type": "start", "messageId": "m1"},
|
||||
{"type": "text-start", "id": "t1"},
|
||||
{"type": "text-delta", "id": "t1", "delta": "Answer "},
|
||||
{"type": "text-delta", "id": "t1", "delta": "is "},
|
||||
{"type": "text-delta", "id": "t1", "delta": "B [citation:42]."},
|
||||
{"type": "text-end", "id": "t1"},
|
||||
{"type": "finish"},
|
||||
])
|
||||
body = _sse_body(
|
||||
[
|
||||
{"type": "start", "messageId": "m1"},
|
||||
{"type": "text-start", "id": "t1"},
|
||||
{"type": "text-delta", "id": "t1", "delta": "Answer "},
|
||||
{"type": "text-delta", "id": "t1", "delta": "is "},
|
||||
{"type": "text-delta", "id": "t1", "delta": "B [citation:42]."},
|
||||
{"type": "text-end", "id": "t1"},
|
||||
{"type": "finish"},
|
||||
]
|
||||
)
|
||||
respx_mock.post("/api/v1/new_chat").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
|
|
@ -208,9 +230,7 @@ async def test_ask_accumulates_text_deltas(respx_mock, http):
|
|||
)
|
||||
)
|
||||
client = NewChatClient(http, _BASE)
|
||||
answer = await client.ask(
|
||||
thread_id=1, search_space_id=2, user_query="What is the answer?"
|
||||
)
|
||||
answer = await client.ask(thread_id=1, search_space_id=2, user_query="What is the answer?")
|
||||
assert answer.text == "Answer is B [citation:42]."
|
||||
assert answer.finished_normally is True
|
||||
assert any(c["chunk_id"] == 42 for c in answer.citations)
|
||||
|
|
@ -219,23 +239,21 @@ async def test_ask_accumulates_text_deltas(respx_mock, http):
|
|||
@pytest.mark.asyncio
|
||||
@respx.mock(base_url=_BASE)
|
||||
async def test_ask_409_thread_busy_retries(respx_mock, http):
|
||||
body = _sse_body([
|
||||
{"type": "text-delta", "id": "t1", "delta": "ok"},
|
||||
{"type": "finish"},
|
||||
])
|
||||
body = _sse_body(
|
||||
[
|
||||
{"type": "text-delta", "id": "t1", "delta": "ok"},
|
||||
{"type": "finish"},
|
||||
]
|
||||
)
|
||||
busy = httpx.Response(
|
||||
409,
|
||||
json={"detail": {"errorCode": "THREAD_BUSY", "message": "busy"}},
|
||||
headers={"Retry-After": "1"},
|
||||
)
|
||||
success = httpx.Response(
|
||||
200, content=body, headers={"Content-Type": "text/event-stream"}
|
||||
)
|
||||
success = httpx.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
|
||||
respx_mock.post("/api/v1/new_chat").mock(side_effect=[busy, success])
|
||||
client = NewChatClient(http, _BASE)
|
||||
answer = await client.ask(
|
||||
thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=2
|
||||
)
|
||||
answer = await client.ask(thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=2)
|
||||
assert answer.text == "ok"
|
||||
|
||||
|
||||
|
|
@ -250,6 +268,4 @@ async def test_ask_409_exhausts_retries(respx_mock, http):
|
|||
respx_mock.post("/api/v1/new_chat").mock(return_value=busy)
|
||||
client = NewChatClient(http, _BASE)
|
||||
with pytest.raises(ThreadBusyError):
|
||||
await client.ask(
|
||||
thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=1
|
||||
)
|
||||
await client.ask(thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=1)
|
||||
|
|
|
|||
|
|
@ -46,32 +46,24 @@ class TestMerge:
|
|||
|
||||
def test_explicit_false_overrides_default_true(self) -> None:
|
||||
defaults = IngestSettings(use_vision_llm=True)
|
||||
merged = IngestSettings.merge(
|
||||
defaults, {"use_vision_llm": False}
|
||||
)
|
||||
merged = IngestSettings.merge(defaults, {"use_vision_llm": False})
|
||||
assert merged.use_vision_llm is False
|
||||
|
||||
def test_explicit_true_overrides_default_false(self) -> None:
|
||||
defaults = IngestSettings(use_vision_llm=False)
|
||||
merged = IngestSettings.merge(
|
||||
defaults, {"use_vision_llm": True}
|
||||
)
|
||||
merged = IngestSettings.merge(defaults, {"use_vision_llm": True})
|
||||
assert merged.use_vision_llm is True
|
||||
|
||||
def test_none_means_silent(self) -> None:
|
||||
# Argparse with BooleanOptionalAction yields None when the
|
||||
# operator passed neither --use-vision-llm nor --no-vision-llm.
|
||||
defaults = IngestSettings(use_vision_llm=True)
|
||||
merged = IngestSettings.merge(
|
||||
defaults, {"use_vision_llm": None}
|
||||
)
|
||||
merged = IngestSettings.merge(defaults, {"use_vision_llm": None})
|
||||
assert merged.use_vision_llm is True
|
||||
|
||||
def test_processing_mode_override(self) -> None:
|
||||
defaults = IngestSettings(processing_mode="basic")
|
||||
merged = IngestSettings.merge(
|
||||
defaults, {"processing_mode": "premium"}
|
||||
)
|
||||
merged = IngestSettings.merge(defaults, {"processing_mode": "premium"})
|
||||
assert merged.processing_mode == "premium"
|
||||
|
||||
def test_processing_mode_invalid_raises(self) -> None:
|
||||
|
|
@ -134,9 +126,7 @@ class TestAddArgs:
|
|||
p = argparse.ArgumentParser()
|
||||
add_ingest_settings_args(
|
||||
p,
|
||||
defaults=IngestSettings(
|
||||
use_vision_llm=False, processing_mode="basic"
|
||||
),
|
||||
defaults=IngestSettings(use_vision_llm=False, processing_mode="basic"),
|
||||
)
|
||||
return p
|
||||
|
||||
|
|
@ -158,31 +148,21 @@ class TestAddArgs:
|
|||
args = parser.parse_args(["--processing-mode", mode])
|
||||
assert args.processing_mode == mode
|
||||
|
||||
def test_processing_mode_rejects_unknown(
|
||||
self, parser: argparse.ArgumentParser
|
||||
) -> None:
|
||||
def test_processing_mode_rejects_unknown(self, parser: argparse.ArgumentParser) -> None:
|
||||
with pytest.raises(SystemExit):
|
||||
parser.parse_args(["--processing-mode", "exotic"])
|
||||
|
||||
def test_vision_flags_mutually_exclusive(
|
||||
self, parser: argparse.ArgumentParser
|
||||
) -> None:
|
||||
def test_vision_flags_mutually_exclusive(self, parser: argparse.ArgumentParser) -> None:
|
||||
with pytest.raises(SystemExit):
|
||||
parser.parse_args(["--use-vision-llm", "--no-vision-llm"])
|
||||
|
||||
def test_full_pipeline(self, parser: argparse.ArgumentParser) -> None:
|
||||
# Operator passes flags + defaults are reasonable. Merge
|
||||
# should yield exactly what they asked for.
|
||||
args = parser.parse_args(
|
||||
["--use-vision-llm", "--processing-mode", "premium"]
|
||||
)
|
||||
defaults = IngestSettings(
|
||||
use_vision_llm=False, processing_mode="basic"
|
||||
)
|
||||
args = parser.parse_args(["--use-vision-llm", "--processing-mode", "premium"])
|
||||
defaults = IngestSettings(use_vision_llm=False, processing_mode="basic")
|
||||
merged = IngestSettings.merge(defaults, vars(args))
|
||||
assert merged == IngestSettings(
|
||||
use_vision_llm=True, processing_mode="premium"
|
||||
)
|
||||
assert merged == IngestSettings(use_vision_llm=True, processing_mode="premium")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -240,16 +220,12 @@ class TestHeader:
|
|||
|
||||
class TestFormatMd:
|
||||
def test_full_settings(self) -> None:
|
||||
out = format_ingest_settings_md(
|
||||
{"use_vision_llm": True, "processing_mode": "premium"}
|
||||
)
|
||||
out = format_ingest_settings_md({"use_vision_llm": True, "processing_mode": "premium"})
|
||||
assert "vision_llm=`on`" in out
|
||||
assert "processing_mode=`premium`" in out
|
||||
|
||||
def test_default_off(self) -> None:
|
||||
out = format_ingest_settings_md(
|
||||
{"use_vision_llm": False, "processing_mode": "basic"}
|
||||
)
|
||||
out = format_ingest_settings_md({"use_vision_llm": False, "processing_mode": "basic"})
|
||||
assert "vision_llm=`off`" in out
|
||||
assert "processing_mode=`basic`" in out
|
||||
|
||||
|
|
|
|||
|
|
@ -25,7 +25,12 @@ from surfsense_evals.core.metrics import (
|
|||
@pytest.mark.parametrize(
|
||||
"k,n,low,high",
|
||||
[
|
||||
(80, 100, 0.7111, 0.8666), # cross-checked vs statsmodels.proportion_confint(method='wilson')
|
||||
(
|
||||
80,
|
||||
100,
|
||||
0.7111,
|
||||
0.8666,
|
||||
), # cross-checked vs statsmodels.proportion_confint(method='wilson')
|
||||
(50, 100, 0.4038, 0.5962),
|
||||
(0, 0, 0.0, 1.0),
|
||||
(0, 10, 0.0, 0.2775),
|
||||
|
|
@ -74,7 +79,7 @@ def test_mcnemar_exact_branch_strong_signal():
|
|||
assert res.b == 0
|
||||
assert res.c == 10
|
||||
assert res.method == "exact"
|
||||
expected = 2 * (0.5 ** 10)
|
||||
expected = 2 * (0.5**10)
|
||||
assert math.isclose(res.p_value, expected, rel_tol=1e-9)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,11 @@ from surfsense_evals.core.parse.answer_letter import AnswerLetterResult
|
|||
@pytest.mark.parametrize(
|
||||
"text,expected_letter,expected_strategy",
|
||||
[
|
||||
('```json\n{"step_by_step_thinking": "...", "answer_choice": "B"}\n```', "B", "json_envelope"),
|
||||
(
|
||||
'```json\n{"step_by_step_thinking": "...", "answer_choice": "B"}\n```',
|
||||
"B",
|
||||
"json_envelope",
|
||||
),
|
||||
('Reasoning... {"step_by_step_thinking": "x", "answer_choice": "C"}', "C", "json_envelope"),
|
||||
("Long reasoning.\nAnswer: D", "D", "answer_line"),
|
||||
("The correct answer is (A).", "A", "answer_line"),
|
||||
|
|
|
|||
|
|
@ -91,7 +91,7 @@ def test_regex_pattern_matches_ts_source():
|
|||
assert "https?://" in pattern
|
||||
assert "urlcite" in pattern
|
||||
assert "doc-" in pattern
|
||||
assert "\u200B" in pattern
|
||||
assert "\u200b" in pattern
|
||||
assert "【" in pattern and "】" in pattern
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -44,11 +44,14 @@ class TestExtractFreeformAnswer:
|
|||
assert extract_freeform_answer("ANSWER: yes") == "yes"
|
||||
assert extract_freeform_answer("answer: no") == "no"
|
||||
|
||||
@pytest.mark.parametrize("text,expected", [
|
||||
("Answer: 1, 2, 3", "1, 2, 3"),
|
||||
("Answer: 3.14", "3.14"),
|
||||
("Answer: spaced ", "spaced"),
|
||||
])
|
||||
@pytest.mark.parametrize(
|
||||
"text,expected",
|
||||
[
|
||||
("Answer: 1, 2, 3", "1, 2, 3"),
|
||||
("Answer: 3.14", "3.14"),
|
||||
("Answer: spaced ", "spaced"),
|
||||
],
|
||||
)
|
||||
def test_various_payloads(self, text: str, expected: str) -> None:
|
||||
assert extract_freeform_answer(text) == expected
|
||||
|
||||
|
|
|
|||
|
|
@ -22,12 +22,16 @@ async def _astream(lines):
|
|||
@pytest.mark.asyncio
|
||||
async def test_basic_data_frame():
|
||||
events = await _alist(
|
||||
iter_sse_events(_astream([
|
||||
'data: {"type": "text-delta", "delta": "hi"}',
|
||||
"",
|
||||
'data: {"type": "finish"}',
|
||||
"",
|
||||
]))
|
||||
iter_sse_events(
|
||||
_astream(
|
||||
[
|
||||
'data: {"type": "text-delta", "delta": "hi"}',
|
||||
"",
|
||||
'data: {"type": "finish"}',
|
||||
"",
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
assert [e.data for e in events] == [
|
||||
'{"type": "text-delta", "delta": "hi"}',
|
||||
|
|
@ -38,10 +42,14 @@ async def test_basic_data_frame():
|
|||
@pytest.mark.asyncio
|
||||
async def test_done_sentinel_passes_through():
|
||||
events = await _alist(
|
||||
iter_sse_events(_astream([
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]))
|
||||
iter_sse_events(
|
||||
_astream(
|
||||
[
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
assert [e.data for e in events] == ["[DONE]"]
|
||||
|
||||
|
|
@ -49,11 +57,15 @@ async def test_done_sentinel_passes_through():
|
|||
@pytest.mark.asyncio
|
||||
async def test_multiline_data_joins_with_newline():
|
||||
events = await _alist(
|
||||
iter_sse_events(_astream([
|
||||
"data: line1",
|
||||
"data: line2",
|
||||
"",
|
||||
]))
|
||||
iter_sse_events(
|
||||
_astream(
|
||||
[
|
||||
"data: line1",
|
||||
"data: line2",
|
||||
"",
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
assert events[0].data == "line1\nline2"
|
||||
|
||||
|
|
@ -61,13 +73,17 @@ async def test_multiline_data_joins_with_newline():
|
|||
@pytest.mark.asyncio
|
||||
async def test_comments_and_other_fields_ignored():
|
||||
events = await _alist(
|
||||
iter_sse_events(_astream([
|
||||
": heartbeat",
|
||||
"event: foo",
|
||||
"id: 123",
|
||||
"data: payload",
|
||||
"",
|
||||
]))
|
||||
iter_sse_events(
|
||||
_astream(
|
||||
[
|
||||
": heartbeat",
|
||||
"event: foo",
|
||||
"id: 123",
|
||||
"data: payload",
|
||||
"",
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
assert [e.data for e in events] == ["payload"]
|
||||
|
||||
|
|
@ -77,8 +93,12 @@ async def test_handles_missing_trailing_blank():
|
|||
"""Some servers omit the final blank line; the consumer should still emit."""
|
||||
|
||||
events = await _alist(
|
||||
iter_sse_events(_astream([
|
||||
"data: only-one",
|
||||
]))
|
||||
iter_sse_events(
|
||||
_astream(
|
||||
[
|
||||
"data: only-one",
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
assert [e.data for e in events] == ["only-one"]
|
||||
|
|
|
|||
|
|
@ -36,11 +36,18 @@ async def test_payload_shape_matches_openrouter_docs(respx_mock, tiny_pdf: Path)
|
|||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"choices": [{
|
||||
"message": {"content": "Answer: B"},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.0001},
|
||||
"choices": [
|
||||
{
|
||||
"message": {"content": "Answer: B"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
"cost": 0.0001,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -63,8 +70,7 @@ async def test_payload_shape_matches_openrouter_docs(respx_mock, tiny_pdf: Path)
|
|||
assert file_part["file"]["filename"] == tiny_pdf.name
|
||||
assert file_part["file"]["file_data"].startswith("data:application/pdf;base64,")
|
||||
assert (
|
||||
base64.b64decode(file_part["file"]["file_data"].split(",", 1)[1])
|
||||
== tiny_pdf.read_bytes() # noqa: ASYNC240 — test fixture, sync read is fine
|
||||
base64.b64decode(file_part["file"]["file_data"].split(",", 1)[1]) == tiny_pdf.read_bytes() # noqa: ASYNC240 — test fixture, sync read is fine
|
||||
)
|
||||
assert user["content"][1] == {"type": "text", "text": "What is the diagnosis?"}
|
||||
assert captured["headers"]["authorization"] == "Bearer sk-or-test"
|
||||
|
|
@ -85,22 +91,22 @@ async def test_chat_array_content_concatenates(respx_mock, tiny_pdf: Path):
|
|||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"choices": [{
|
||||
"message": {
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello "},
|
||||
{"type": "text", "text": "world"},
|
||||
{"type": "image_url", "image_url": "ignored"},
|
||||
]
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello "},
|
||||
{"type": "text", "text": "world"},
|
||||
{"type": "image_url", "image_url": "ignored"},
|
||||
]
|
||||
}
|
||||
}
|
||||
}],
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
|
||||
},
|
||||
)
|
||||
)
|
||||
provider = OpenRouterPdfProvider(
|
||||
api_key="sk-or-test", base_url=_BASE, model="x/y"
|
||||
)
|
||||
provider = OpenRouterPdfProvider(api_key="sk-or-test", base_url=_BASE, model="x/y")
|
||||
response = await provider.complete(prompt="hi", pdf_path=tiny_pdf)
|
||||
assert response.text == "Hello world"
|
||||
|
||||
|
|
|
|||
|
|
@ -65,13 +65,15 @@ class TestParser:
|
|||
interaction_id="abc",
|
||||
query="Who directed Inception?",
|
||||
answer="Christopher Nolan",
|
||||
pages=[{
|
||||
"page_name": "Inception (film)",
|
||||
"page_url": "https://en.wikipedia.org/wiki/Inception",
|
||||
"page_snippet": "snippet",
|
||||
"page_result": "<html>full html</html>",
|
||||
"page_last_modified": "2024-01-01",
|
||||
}],
|
||||
pages=[
|
||||
{
|
||||
"page_name": "Inception (film)",
|
||||
"page_url": "https://en.wikipedia.org/wiki/Inception",
|
||||
"page_snippet": "snippet",
|
||||
"page_result": "<html>full html</html>",
|
||||
"page_last_modified": "2024-01-01",
|
||||
}
|
||||
],
|
||||
),
|
||||
]
|
||||
path = _make_jsonl_bz2(rows, tmp_path)
|
||||
|
|
@ -120,8 +122,7 @@ class TestParser:
|
|||
|
||||
def test_alt_answers_parsed(self, tmp_path: Path) -> None:
|
||||
rows = [
|
||||
_row(interaction_id="z", query="q?", answer="42",
|
||||
alt_ans=["forty-two", "42.0"]),
|
||||
_row(interaction_id="z", query="q?", answer="42", alt_ans=["forty-two", "42.0"]),
|
||||
]
|
||||
path = _make_jsonl_bz2(rows, tmp_path)
|
||||
parsed = iter_questions(path)
|
||||
|
|
@ -143,22 +144,32 @@ class TestParser:
|
|||
class TestPageHash:
|
||||
def test_url_hash_stable(self) -> None:
|
||||
a = CragPage(
|
||||
page_name="A", page_url="https://x.test/p?q=1",
|
||||
page_snippet="", page_html="<html/>",
|
||||
page_name="A",
|
||||
page_url="https://x.test/p?q=1",
|
||||
page_snippet="",
|
||||
page_html="<html/>",
|
||||
)
|
||||
b = CragPage(
|
||||
page_name="B", page_url="https://x.test/p?q=1",
|
||||
page_snippet="", page_html="<html/>",
|
||||
page_name="B",
|
||||
page_url="https://x.test/p?q=1",
|
||||
page_snippet="",
|
||||
page_html="<html/>",
|
||||
)
|
||||
assert a.url_hash == b.url_hash
|
||||
assert len(a.url_hash) == 12
|
||||
|
||||
def test_url_hash_unique(self) -> None:
|
||||
a = CragPage(
|
||||
page_name="A", page_url="https://x.test/a", page_snippet="", page_html="<html/>",
|
||||
page_name="A",
|
||||
page_url="https://x.test/a",
|
||||
page_snippet="",
|
||||
page_html="<html/>",
|
||||
)
|
||||
b = CragPage(
|
||||
page_name="B", page_url="https://x.test/b", page_snippet="", page_html="<html/>",
|
||||
page_name="B",
|
||||
page_url="https://x.test/b",
|
||||
page_snippet="",
|
||||
page_html="<html/>",
|
||||
)
|
||||
assert a.url_hash != b.url_hash
|
||||
|
||||
|
|
@ -174,21 +185,23 @@ class TestStratifiedSample:
|
|||
(5, "sports", "multi-hop"),
|
||||
):
|
||||
for _ in range(n):
|
||||
out.append(CragQuestion(
|
||||
qid=f"C{idx:05d}",
|
||||
interaction_id=f"i{idx}",
|
||||
query_time="2024-01-01",
|
||||
query=f"q{idx}?",
|
||||
gold_answer="a",
|
||||
alt_answers=[],
|
||||
domain=domain,
|
||||
question_type=qtype,
|
||||
static_or_dynamic="static",
|
||||
popularity="head",
|
||||
split=0,
|
||||
raw_index=idx,
|
||||
pages=[],
|
||||
))
|
||||
out.append(
|
||||
CragQuestion(
|
||||
qid=f"C{idx:05d}",
|
||||
interaction_id=f"i{idx}",
|
||||
query_time="2024-01-01",
|
||||
query=f"q{idx}?",
|
||||
gold_answer="a",
|
||||
alt_answers=[],
|
||||
domain=domain,
|
||||
question_type=qtype,
|
||||
static_or_dynamic="static",
|
||||
popularity="head",
|
||||
split=0,
|
||||
raw_index=idx,
|
||||
pages=[],
|
||||
)
|
||||
)
|
||||
idx += 1
|
||||
return out
|
||||
|
||||
|
|
|
|||
|
|
@ -152,7 +152,9 @@ class TestGradeDeterministicHappyPath:
|
|||
class TestGradeDeterministicRefusal:
|
||||
def test_idk_maps_to_missing(self) -> None:
|
||||
result = grade_deterministic(
|
||||
pred="I don't know.", gold="Tim Cook", question_type="simple",
|
||||
pred="I don't know.",
|
||||
gold="Tim Cook",
|
||||
question_type="simple",
|
||||
)
|
||||
assert result.grade == "missing"
|
||||
assert result.score == 0
|
||||
|
|
@ -225,8 +227,11 @@ class TestGradeDeterministicLexicalMiss:
|
|||
class TestGradeResultShape:
|
||||
def test_to_dict_round_trip(self) -> None:
|
||||
result = CragGradeResult(
|
||||
grade="correct", score=1, method="exact",
|
||||
normalised_pred="x", normalised_gold="x",
|
||||
grade="correct",
|
||||
score=1,
|
||||
method="exact",
|
||||
normalised_pred="x",
|
||||
normalised_gold="x",
|
||||
)
|
||||
d = result.to_dict()
|
||||
assert d["grade"] == "correct"
|
||||
|
|
|
|||
|
|
@ -112,7 +112,9 @@ class TestFallbackStripper:
|
|||
</body></html>
|
||||
"""
|
||||
result = extract_main_content(
|
||||
html, url="https://x.test/", page_name="Title",
|
||||
html,
|
||||
url="https://x.test/",
|
||||
page_name="Title",
|
||||
)
|
||||
assert result.ok
|
||||
assert "content one" in result.text
|
||||
|
|
|
|||
|
|
@ -63,14 +63,22 @@ class TestCacheFilename:
|
|||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_fetch_success_writes_markdown(tmp_path: Path) -> None:
|
||||
respx.get(WIKI_API).mock(return_value=httpx.Response(
|
||||
200,
|
||||
json={"query": {"pages": [{
|
||||
"pageid": 1,
|
||||
"title": "James Buchanan",
|
||||
"extract": "James Buchanan was the 15th president of the United States.",
|
||||
}]}},
|
||||
))
|
||||
respx.get(WIKI_API).mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"query": {
|
||||
"pages": [
|
||||
{
|
||||
"pageid": 1,
|
||||
"title": "James Buchanan",
|
||||
"extract": "James Buchanan was the 15th president of the United States.",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
)
|
||||
fetcher = WikiFetcher(cache_dir=tmp_path, rate_limit_rps=100) # disable throttle
|
||||
article = await fetcher.fetch("https://en.wikipedia.org/wiki/James_Buchanan")
|
||||
assert article is not None
|
||||
|
|
@ -83,13 +91,21 @@ async def test_fetch_success_writes_markdown(tmp_path: Path) -> None:
|
|||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_fetch_missing_page_returns_none(tmp_path: Path) -> None:
|
||||
respx.get(WIKI_API).mock(return_value=httpx.Response(
|
||||
200,
|
||||
json={"query": {"pages": [{
|
||||
"title": "DoesNotExist",
|
||||
"missing": True,
|
||||
}]}},
|
||||
))
|
||||
respx.get(WIKI_API).mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"query": {
|
||||
"pages": [
|
||||
{
|
||||
"title": "DoesNotExist",
|
||||
"missing": True,
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
)
|
||||
fetcher = WikiFetcher(cache_dir=tmp_path, rate_limit_rps=100)
|
||||
article = await fetcher.fetch("https://en.wikipedia.org/wiki/DoesNotExist")
|
||||
assert article is None
|
||||
|
|
|
|||
|
|
@ -99,7 +99,9 @@ class TestListFormat:
|
|||
assert 0.0 < r.f1 < 1.0
|
||||
|
||||
def test_extra_items_lower_precision(self) -> None:
|
||||
r = grade(pred="apple, banana, cherry, date", gold="apple, banana, cherry", answer_format="List")
|
||||
r = grade(
|
||||
pred="apple, banana, cherry, date", gold="apple, banana, cherry", answer_format="List"
|
||||
)
|
||||
assert 0.0 < r.f1 < 1.0
|
||||
# Recall=1, precision=3/4 → F1 ~= 0.857
|
||||
assert r.f1 == pytest.approx(2 * (3 / 4) * 1 / (3 / 4 + 1), rel=1e-3)
|
||||
|
|
|
|||
|
|
@ -14,9 +14,7 @@ from ...core.workspace_context import WorkspaceContext
|
|||
from . import document_tools, search_tools
|
||||
|
||||
|
||||
def register(
|
||||
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
|
||||
) -> None:
|
||||
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
|
||||
"""Register every knowledge-base tool on the server."""
|
||||
search_tools.register(mcp, client, context)
|
||||
document_tools.register(mcp, client, context)
|
||||
|
|
|
|||
|
|
@ -20,9 +20,7 @@ from .annotations import DELETE, WRITE, DocumentId
|
|||
from .note_ingestion import build_note_document
|
||||
|
||||
|
||||
def register(
|
||||
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
|
||||
) -> None:
|
||||
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
|
||||
"""Register the knowledge-base write and delete tools."""
|
||||
|
||||
@mcp.tool(
|
||||
|
|
@ -136,8 +134,7 @@ def register(
|
|||
str,
|
||||
Field(
|
||||
min_length=1,
|
||||
description="New full text; replaces the existing content "
|
||||
"entirely.",
|
||||
description="New full text; replaces the existing content entirely.",
|
||||
),
|
||||
],
|
||||
) -> str:
|
||||
|
|
|
|||
|
|
@ -17,9 +17,7 @@ from ...core.workspace_context import WorkspaceContext, WorkspaceParam
|
|||
from .annotations import READ, DocumentId, DocumentTypes
|
||||
|
||||
|
||||
def register(
|
||||
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
|
||||
) -> None:
|
||||
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
|
||||
"""Register the knowledge-base read tools."""
|
||||
|
||||
@mcp.tool(
|
||||
|
|
@ -81,12 +79,8 @@ def register(
|
|||
int | None,
|
||||
Field(description="Only documents in this folder. Omit for all."),
|
||||
] = None,
|
||||
page: Annotated[
|
||||
int, Field(ge=0, description="Zero-based page number.")
|
||||
] = 0,
|
||||
page_size: Annotated[
|
||||
int, Field(ge=1, description="Documents per page.")
|
||||
] = 20,
|
||||
page: Annotated[int, Field(ge=0, description="Zero-based page number.")] = 0,
|
||||
page_size: Annotated[int, Field(ge=1, description="Documents per page.")] = 20,
|
||||
workspace: WorkspaceParam = None,
|
||||
response_format: ResponseFormatParam = "markdown",
|
||||
) -> str:
|
||||
|
|
|
|||
|
|
@ -35,9 +35,7 @@ _REGISTRARS = (
|
|||
)
|
||||
|
||||
|
||||
def register(
|
||||
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
|
||||
) -> None:
|
||||
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
|
||||
"""Register every scraper and run-history tool on the server."""
|
||||
for module in _REGISTRARS:
|
||||
module.register(mcp, client, context)
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue