fund_rfid_data/llm_extract.py
Florian Herzog 09798eb27a Add LLM grounding pipeline: current-source fetch, alias + LLM role-check matching
Ensures text<->gold agreement for the text->triple dataset:
- fetch: newest full prospectus book + later 497 supplements (time-aligned with
  the current N-CEN gold; fixes stale 'largest book' picking 2-5yr-old filings)
- grounding: fast alias matcher (name present, variant-tolerant) AND an LLM
  role-check (llm_extract.py match mode, via local Ollama or remote vLLM server)
  that verifies the entity plays that ROLE -- catches right-name/wrong-role cases
  a lexical matcher over-keeps. Validated with a strong model; ~93% of gold names
  are present once granularity+currency are fixed.
- llm_extract.py: extract (baseline) + match (grounding) modes, sliding-window
  with retrieval pre-filter, claim dedup, retry, ollama/vllm backends
- build_rdf_dataset.py: --grounding alias|llm|name|context|none, whole-trust
  samples now filtered, MIN_PROSE stub guard

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-09 13:45:32 +02:00

495 lines
21 KiB
Python

#!/usr/bin/env python3
"""
LLM extraction baseline for the SEC prospectus -> RDF triple dataset.
A strong open-source instruct model (default qwen3.6:35b via local Ollama) reads
each fund's prospectus segment together with the target ontology and extracts the
service-provider triples. This serves two purposes for the thesis:
1. STRONG BASELINE: the prompted model's triples, scored against the N-CEN gold
with score_baseline.py, measure how well a state-of-the-art model does the
task without finetuning -- the comparison point for the four finetuned models.
2. SEMANTIC GROUNDING: whether a gold triple is actually extractable from the
text is a semantic question that lexical substring/keyword matching cannot
answer (e.g. "John Hancock" matches the fund heading, not the administrator
role; "administration services are provided by X" is a valid mention no fixed
keyword catches). What the strong model extracts from the segment is therefore
a far better grounding signal than the lexical flag.
The model is asked to extract ONLY relations present in the supplied ontology and
to copy entity names verbatim from the text, so that a fair triple-level
comparison against the gold is possible.
Usage:
# all (or a sample of) the dataset
python llm_extract.py --in data/rdf_poc/test.jsonl --out data/rdf_poc/preds_qwen.jsonl
python llm_extract.py --in data/rdf_poc/samples.jsonl --limit 50 --model qwen3.6:35b
# then score:
python score_baseline.py model --pred data/rdf_poc/preds_qwen.jsonl
"""
import argparse
import json
import re
import time
import urllib.request
from pathlib import Path
OLLAMA_URL = "http://localhost:11434/api/chat"
# Backend config (set in main()). Two backends are supported:
# "ollama" -> local Ollama (/api/chat), per-request num_ctx
# "vllm" -> remote vLLM OpenAI-compatible server (/v1/chat/completions),
# fixed served context; much higher throughput / stronger model
# (e.g. silicon.fhgr.ch:7080 serving qwen3.5-122b).
BACKEND = "ollama"
VLLM_URL = "http://silicon.fhgr.ch:7080/v1/chat/completions"
API_KEY = ""
SYSTEM = (
"You are an information-extraction engine. You read a U.S. mutual-fund "
"prospectus/SAI excerpt and extract relationships between named entities as "
"RDF-style triples. You output ONLY valid JSON, no commentary."
)
# The instruction is deliberately strict: only ontology relations, names copied
# verbatim, and an explicit 'not stated' contract so the model does NOT invent
# service providers that the text does not actually name (the whole point of the
# grounding test).
PROMPT_TMPL = """\
Extract relationships for the fund described in the TEXT below.
Allowed relations (subject -> predicate -> object type), from the ontology:
{ontology}
Rules:
- Extract a triple ONLY if the TEXT explicitly states that relationship for this
fund. Do NOT use outside knowledge. If a relation is not stated, omit it.
- Copy each entity name VERBATIM as it appears in the TEXT (full legal name).
- Use only the predicates listed above.
- Output JSON of the form:
{{"triples": [{{"subject": "...", "predicate": "...", "object": "..."}}, ...]}}
TEXT:
\"\"\"
{text}
\"\"\"
"""
MAX_TEXT_CHARS = 1_200_000 # ~300k tokens; effectively the whole single book
# (set lower only to bound latency; too-low truncates
# the SAI section where service providers are named)
def ollama_chat(model, system, prompt, timeout=600, num_ctx=32768):
"""Send a chat request to the active backend; return the message content.
Name kept for backward compatibility with existing call sites.
"""
messages = [{"role": "system", "content": system},
{"role": "user", "content": prompt}]
if BACKEND == "vllm":
body = {
"model": model,
"messages": messages,
"temperature": 0,
# Reasoning model: it spends several thousand tokens "thinking" before
# emitting JSON. Too small a budget -> it never reaches the answer and
# returns empty content. Keep this generous.
"max_tokens": 8000,
"response_format": {"type": "json_object"},
}
headers = {"Content-Type": "application/json"}
if API_KEY:
headers["Authorization"] = f"Bearer {API_KEY}"
req = urllib.request.Request(
VLLM_URL, data=json.dumps(body).encode("utf-8"), headers=headers)
with urllib.request.urlopen(req, timeout=timeout) as r:
resp = json.loads(r.read().decode("utf-8", "replace"))
# reasoning models can return content=null (spent budget "thinking");
# fall back to an empty string so callers' parsers degrade gracefully.
return resp["choices"][0]["message"].get("content") or ""
# default: ollama
body = {
"model": model,
"messages": messages,
"stream": False,
"format": "json", # ask Ollama to constrain to JSON
"options": {"temperature": 0, "num_ctx": num_ctx},
}
req = urllib.request.Request(
OLLAMA_URL, data=json.dumps(body).encode("utf-8"),
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=timeout) as r:
resp = json.loads(r.read().decode("utf-8", "replace"))
return resp["message"]["content"]
def parse_triples(raw):
"""Parse the model's JSON; tolerate stray text/code-fences around it."""
s = raw.strip()
s = re.sub(r"^```(?:json)?|```$", "", s, flags=re.MULTILINE).strip()
try:
obj = json.loads(s)
except Exception:
m = re.search(r"\{.*\}", s, re.S)
if not m:
return []
try:
obj = json.loads(m.group(0))
except Exception:
return []
out = []
for t in obj.get("triples", []):
s_, p_, o_ = t.get("subject"), t.get("predicate"), t.get("object")
if s_ and p_ and o_:
out.append({"s": s_, "p": p_, "o": o_})
return out
def ontology_str(ont):
lines = []
for st, preds in ont.items():
for p, ots in preds.items():
lines.append(f" {st} -> {p} -> {', '.join(ots)}")
return "\n".join(lines)
def _dedup(triples):
seen, out = set(), []
for t in triples:
k = (t["p"].lower(), t["o"].strip().lower())
if k not in seen:
seen.add(k)
out.append(t)
return out
def extract_windowed(model, ontology, text, win_chars, overlap, num_ctx, timeout=600):
"""Slide a window over a long document, extract per window, union the results.
A single huge-context call over a 500KB+ book is slow and suffers from
'lost in the middle'; instead we pass overlapping windows (each small enough
for fast inference and a modest context) and take the de-duplicated union of
the triples found in each. Service-provider facts that live only in the SAI
section near the end are thus reliably seen by some window. Overlap keeps a
provider sentence that straddles a window boundary intact in at least one
window.
"""
step = max(1, win_chars - overlap)
windows = [text[i:i + win_chars] for i in range(0, max(1, len(text)), step)]
all_triples, n_win = [], 0
for w in windows:
if len(w.strip()) < 200:
continue
prompt = PROMPT_TMPL.format(ontology=ontology, text=w)
ctx = min(num_ctx, max(8192, int(len(prompt) / 3.2) + 2048))
try:
raw = ollama_chat(model, SYSTEM, prompt, timeout=timeout, num_ctx=ctx)
all_triples.extend(parse_triples(raw))
except Exception as e:
print(f" window error: {e}")
n_win += 1
return _dedup(all_triples), n_win
# --------------------------------------------------------------------------- #
# match mode: ask the LLM whether each GOLD triple is supported by the text.
# This sidesteps the surface-form problem entirely: the model decides semantic
# equivalence (e.g. "BNY Mellon" == "The Bank of New York Mellon", "the
# Administrator" anaphora) which no fuzzy string rule can capture.
# --------------------------------------------------------------------------- #
MATCH_SYSTEM = (
"You verify facts against a source text. Given a U.S. mutual-fund "
"prospectus/SAI excerpt and a list of candidate (entity, role, fund) claims, "
"you decide for EACH claim whether the excerpt supports it. Treat name "
"variants, abbreviations and anaphora as the same entity (e.g. 'BNY Mellon' "
"= 'The Bank of New York Mellon'). Output ONLY valid JSON."
)
MATCH_PROMPT = """\
For each numbered CLAIM below, decide whether the TEXT supports it for the fund.
A claim is supported only if the TEXT states that the named entity plays that
role (advisedBy=investment adviser, subAdvisedBy=sub-adviser,
transferAgent=transfer agent, custodian=custodian, administrator=administrator,
underwrittenBy=distributor/underwriter, seriesOf=the fund is a series of that
trust). Use ONLY the TEXT; treat name variants/abbreviations as equal.
Output JSON: {{"results": [{{"id": <n>, "supported": true|false}}, ...]}} with one
entry per claim id.
CLAIMS:
{claims}
TEXT:
\"\"\"
{text}
\"\"\"
"""
def _claim_lines(claims):
return "\n".join(
f' {i}. entity="{c["o"]}" role={c["p"]}' for i, c in enumerate(claims))
def parse_results(raw, n):
"""Parse {"results":[{id,supported}]} -> list[bool] of length n (default False)."""
out = [False] * n
s = re.sub(r"^```(?:json)?|```$", "", raw.strip(), flags=re.MULTILINE).strip()
try:
obj = json.loads(s)
except Exception:
m = re.search(r"\{.*\}", s, re.S)
if not m:
return out
try:
obj = json.loads(m.group(0))
except Exception:
return out
for r in obj.get("results", []):
try:
i = int(r["id"])
except (KeyError, ValueError, TypeError):
continue
if 0 <= i < n:
out[i] = bool(r.get("supported"))
return out
# Relation keywords that mark a window as possibly relevant to service-provider
# claims. A window with none of these AND no gold-name token cannot support any
# claim, so it is skipped (retrieval pre-filter) -- this avoids sending the bulk
# of a multi-MB book (holdings tables, risk boilerplate) to the LLM.
_RELEVANCE_KW = ("custodian", "custody", "adviser", "advisor", "sub-advis",
"transfer agent", "administrat", "distributor", "underwriter",
"serves as", "acts as", "series of", "a series")
_STOPTOK = {"the", "inc", "llc", "lp", "company", "trust", "fund", "funds", "of",
"and", "group", "capital", "asset", "management", "services", "co",
"national", "association", "bank"}
def _name_tokens(label):
return [w for w in re.sub(r"[^a-z0-9 ]", " ", (label or "").lower()).split()
if w not in _STOPTOK and len(w) >= 4]
def _window_relevant(w_low, name_tokens):
if any(k in w_low for k in _RELEVANCE_KW):
return True
return any(tok in w_low for tok in name_tokens)
def match_windowed(model, gold_triples, text, win_chars, overlap, num_ctx, timeout=600):
"""For each gold triple, is it supported somewhere in the text?
Retrieval pre-filter: only windows that contain a relation keyword or a
distinctive token of some gold object name are sent to the LLM (the rest of a
multi-MB book -- holdings tables, risk sections -- cannot support a
service-provider claim). All gold triples are checked together per kept window
(one call/window); a triple counts as grounded if ANY window supports it.
Returns (list[bool], n_windows_sent).
"""
claims = _claim_lines(gold_triples)
name_tokens = set()
for t in gold_triples:
name_tokens.update(_name_tokens(t.get("o", "")))
step = max(1, win_chars - overlap)
windows = [text[i:i + win_chars] for i in range(0, max(1, len(text)), step)]
supported = [False] * len(gold_triples)
n_win = n_fail = 0
for w in windows:
if len(w.strip()) < 200:
continue
if not _window_relevant(w.lower(), name_tokens):
continue # retrieval pre-filter: no provider signal in this window
prompt = MATCH_PROMPT.format(claims=claims, text=w)
ctx = min(num_ctx, max(8192, int(len(prompt) / 3.2) + 2048))
ok = False
for attempt in range(2): # one retry on transient failure
try:
raw = ollama_chat(model, MATCH_SYSTEM, prompt, timeout=timeout, num_ctx=ctx)
res = parse_results(raw, len(gold_triples))
supported = [a or b for a, b in zip(supported, res)]
ok = True
break
except Exception as e:
if attempt == 0:
time.sleep(3)
else:
print(f" window error: {e}")
if not ok:
n_fail += 1
n_win += 1
# early exit: once every claim is supported, no need to scan further
if all(supported):
break
# If windows failed, the 'False' entries are unreliable (could be missed calls,
# not true absences). Signal this so the caller does not record bogus 0s.
return supported, n_win, n_fail
def _match_one(args, r):
"""Match all gold triples of one sample; return the annotated record."""
text = r["input_text"][:MAX_TEXT_CHARS]
gold = r.get("target_triples", [])
if not gold:
return {"sample_id": r.get("sample_id"), "cik": r["cik"],
"fund": r.get("fund"), "n_windows": 0, "triples": []}
# DEDUP claims: all funds of a trust share the same providers, so e.g. 540
# triples collapse to ~13 distinct (predicate, name) pairs. Matching the
# distinct set keeps each LLM call small (a long claim list overflows the
# answer's max_tokens and yields truncated JSON -> all False) and fast; the
# per-distinct verdict is then mapped back to every triple.
distinct, key_of = [], []
seen = {}
for t in gold:
k = (t["p"], (t.get("o") or "").strip().lower())
if k not in seen:
seen[k] = len(distinct)
distinct.append({"p": t["p"], "o": t.get("o")})
key_of.append(seen[k])
n_fail = 0
try:
if args.window and len(text) > args.window:
d_sup, n_win, n_fail = match_windowed(
args.model, distinct, text, args.window, args.overlap, args.num_ctx)
else:
prompt = MATCH_PROMPT.format(claims=_claim_lines(distinct), text=text)
ctx = min(args.num_ctx, max(8192, int(len(prompt) / 3.2) + 2048))
raw = ollama_chat(args.model, MATCH_SYSTEM, prompt, num_ctx=ctx)
d_sup = parse_results(raw, len(distinct)); n_win = 1
except Exception as e:
print(f" [{r.get('sample_id')}] ERROR: {e}")
d_sup = [False] * len(distinct); n_win = 0; n_fail = 1
ann = [{**t, "llm_grounded": bool(d_sup[key_of[i]])} for i, t in enumerate(gold)]
return {"sample_id": r.get("sample_id"), "cik": r["cik"], "fund": r.get("fund"),
"n_windows": n_win, "n_failed_windows": n_fail, "n_distinct": len(distinct),
"ok": n_fail == 0, "triples": ann}
def run_match(args):
"""LLM match mode: annotate each gold triple with llm_grounded (semantic).
Samples are processed concurrently (--workers) to exploit the vLLM server's
parallel request handling (max-num-seqs); each sample's windows still run
sequentially within its worker. Results are written as they complete.
"""
from concurrent.futures import ThreadPoolExecutor, as_completed
rows = [json.loads(l) for l in open(args.inp, encoding="utf-8")]
if args.limit:
rows = rows[:args.limit]
done = tot_g = tot_t = 0
t0 = time.time()
workers = max(1, args.workers)
with open(args.out, "w", encoding="utf-8") as out, \
ThreadPoolExecutor(max_workers=workers) as ex:
futs = {ex.submit(_match_one, args, r): r for r in rows}
for fut in as_completed(futs):
rec = fut.result()
ng = sum(t.get("llm_grounded") for t in rec["triples"])
tot_g += ng; tot_t += len(rec["triples"])
out.write(json.dumps(rec, ensure_ascii=False) + "\n"); out.flush()
done += 1
dt = time.time() - t0
print(f"[{done}/{len(rows)}] {rec.get('fund') or rec['cik']:.38s} "
f"grounded {ng}/{len(rec['triples'])} "
f"({dt/done:.1f}s/sample avg, {workers}w)")
print(f"\nmatched {done} samples -> {args.out} "
f"(LLM-grounded {tot_g}/{tot_t} = {100*tot_g/max(1,tot_t):.0f}%)")
def run_extract(args):
rows = [json.loads(l) for l in open(args.inp, encoding="utf-8")]
if args.limit:
rows = rows[:args.limit]
done = 0
t0 = time.time()
with open(args.out, "w", encoding="utf-8") as out:
for r in rows:
text = r["input_text"][:MAX_TEXT_CHARS]
ont = ontology_str(r["ontology"])
n_win = 1
try:
if args.window and len(text) > args.window:
# long document -> sliding-window extraction + union
triples, n_win = extract_windowed(
args.model, ont, text, args.window, args.overlap, args.num_ctx)
else:
prompt = PROMPT_TMPL.format(ontology=ont, text=text)
ctx = min(args.num_ctx, max(8192, int(len(prompt) / 3.2) + 2048))
triples = _dedup(parse_triples(
ollama_chat(args.model, SYSTEM, prompt, num_ctx=ctx)))
except Exception as e:
print(f" [{r.get('sample_id')}] ERROR: {e}")
triples = []
rec = {"sample_id": r.get("sample_id"), "cik": r["cik"],
"fund": r.get("fund"), "n_windows": n_win, "triples": triples}
out.write(json.dumps(rec, ensure_ascii=False) + "\n")
out.flush()
done += 1
dt = time.time() - t0
wtag = f" [{n_win}w]" if n_win > 1 else ""
print(f"[{done}/{len(rows)}] {r.get('fund') or r['cik']:.40s} "
f"-> {len(triples)} triples{wtag} ({dt/done:.1f}s/sample)")
print(f"\nwrote {done} predictions -> {args.out} ({args.model})")
def main():
global BACKEND, VLLM_URL, API_KEY
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--mode", choices=["extract", "match"], default="extract",
help="extract = free triple extraction (baseline); "
"match = verify each GOLD triple against the text "
"(semantic grounding flag llm_grounded)")
ap.add_argument("--in", dest="inp", default="data/rdf_poc/test.jsonl")
ap.add_argument("--out", dest="out", default="data/rdf_poc/preds_qwen.jsonl")
ap.add_argument("--backend", choices=["ollama", "vllm"], default="ollama",
help="ollama = local Mac; vllm = remote OpenAI-compatible "
"server (stronger model, higher throughput)")
ap.add_argument("--model", default=None,
help="model id; defaults to qwen3.6:35b (ollama) or "
"qwen3.5-122b-a10b-fp8 (vllm)")
ap.add_argument("--vllm-url", default=VLLM_URL)
ap.add_argument("--api-key", default="")
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--num-ctx", type=int, default=32768)
ap.add_argument("--window", type=int, default=0,
help="sliding-window size in chars; texts longer than this are "
"processed window-by-window and unioned. 0 = auto: a "
"server-safe size derived from --num-ctx (vllm) or ~600 KB "
"(ollama, where one big call is faster).")
ap.add_argument("--overlap", type=int, default=8000,
help="overlap between consecutive windows in chars")
ap.add_argument("--workers", type=int, default=1,
help="concurrent samples (vllm server handles parallel "
"requests; keep <= the server's max-num-seqs, e.g. 8-16)")
args = ap.parse_args()
BACKEND = args.backend
VLLM_URL = args.vllm_url
API_KEY = args.api_key
if not args.model:
args.model = "qwen3.5-122b-a10b-fp8" if BACKEND == "vllm" else "qwen3.6:35b"
if not args.window:
if BACKEND == "vllm":
# Empirically (128k-context reasoning server) ~300 KB / ~93k tokens is
# the reliable ceiling per call: above it the model spends its whole
# output budget "thinking" and returns empty (finish_reason=length) or
# times out. 280 KB leaves headroom and yields 1-3 windows for a
# typical book while staying fast (~80s/call).
args.window = 280000
else:
args.window = 600000 # ollama: one big local call beats many windows
if args.mode == "match":
run_match(args)
else:
run_extract(args)
if __name__ == "__main__":
main()