991715ab76
- llm_extract.py: match mode now window-parallel with retrieval pre-filter,
claim dedup, retry, and enable_thinking=false (vLLM) -> ~36x faster per call;
n_failed_windows/ok flags so an interrupted run never records bogus 0s.
- build_rdf_dataset.py:
- gold now includes the share-class level (hasShareClass/ticker/className)
- grounding modes alias|llm|name|context|none (--grounding); llm reads the
role-check verdicts from match_all.jsonl
- label stage: per-triple extractable + per-sample FULL/PARTIAL/NONE
- trainset stage: combines GROUNDED triples with focused TEXT EXCERPTS cut
around the actual provider statement (evidence), not the multi-MB book
- split --src to split trainset.jsonl (trust-level, no leakage)
- helper scripts: watch_match.sh, resume_match.sh (crash/sleep-safe resume),
finalize_dataset.sh
- final dataset: 335/335 trusts, 85% text<->gold agreement, 334 samples,
10,689 grounded triples, train/val/test 264/35/35
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
558 lines
24 KiB
Python
558 lines
24 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 = ""
|
|
NO_THINK = True # disable reasoning "thinking" on vLLM (huge speedup, see below)
|
|
WINDOW_WORKERS = 8 # concurrent windows per trust (trust-workers x this <= server cap)
|
|
|
|
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,
|
|
"max_tokens": 8000,
|
|
"response_format": {"type": "json_object"},
|
|
}
|
|
# Disable the reasoning model's "thinking": for this yes/no grounding task
|
|
# it spends thousands of tokens deliberating before the JSON, making each
|
|
# call ~16x slower (7.8s -> 0.5s measured) with no quality gain.
|
|
if NO_THINK:
|
|
body["chat_template_kwargs"] = {"enable_thinking": False}
|
|
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.
|
|
Use ONLY the TEXT; treat name variants/abbreviations as equal (e.g. "BNY Mellon"
|
|
= "The Bank of New York Mellon").
|
|
|
|
Two kinds of claim:
|
|
- ROLE claims (advisedBy, subAdvisedBy, transferAgent, custodian, administrator,
|
|
underwrittenBy, seriesOf): supported only if the TEXT states the named entity
|
|
plays that role for the fund (advisedBy=investment adviser, underwrittenBy=
|
|
distributor/underwriter, seriesOf=the fund is a series of that trust, etc.).
|
|
- PRESENCE claims (hasShareClass, className, ticker): supported if the value
|
|
simply appears in the TEXT — the share-class name or ticker symbol is printed
|
|
anywhere (e.g. in the fee table or cover). No role needed.
|
|
|
|
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=150):
|
|
"""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)]
|
|
# keep only relevant windows (retrieval pre-filter)
|
|
kept = [w for w in windows
|
|
if len(w.strip()) >= 200 and _window_relevant(w.lower(), name_tokens)]
|
|
supported = [False] * len(gold_triples)
|
|
evidence = [None] * len(gold_triples) # supporting window text per triple
|
|
n_fail = 0
|
|
|
|
def run_window(w):
|
|
prompt = MATCH_PROMPT.format(claims=claims, text=w)
|
|
ctx = min(num_ctx, max(8192, int(len(prompt) / 3.2) + 2048))
|
|
for attempt in range(2):
|
|
try:
|
|
raw = ollama_chat(model, MATCH_SYSTEM, prompt, timeout=timeout, num_ctx=ctx)
|
|
return parse_results(raw, len(gold_triples)), w
|
|
except Exception as e:
|
|
if attempt == 0:
|
|
time.sleep(3)
|
|
else:
|
|
print(f" window error: {e}")
|
|
return None, w
|
|
|
|
# Send a trust's windows CONCURRENTLY: one big book's ~10-15 windows then
|
|
# saturate the server's batch on their own (sequential per-window left the
|
|
# server idle). WINDOW_WORKERS caps in-flight calls per trust.
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
if kept:
|
|
with ThreadPoolExecutor(max_workers=min(WINDOW_WORKERS, len(kept))) as ex:
|
|
for res, w in ex.map(run_window, kept):
|
|
if res is None:
|
|
n_fail += 1
|
|
continue
|
|
for i, s in enumerate(res):
|
|
if s and not supported[i]:
|
|
supported[i] = True
|
|
evidence[i] = w # first window that supports triple i
|
|
return supported, len(kept), n_fail, evidence
|
|
|
|
|
|
def _excerpt(window, name, radius=400):
|
|
"""Return a tight excerpt of `window` around the first occurrence of `name`.
|
|
|
|
Locates the most distinctive token of the entity name and returns +/- radius
|
|
chars around it (the evidence the model used). Falls back to the window head
|
|
if the token is not locatable. Returns None if no window.
|
|
"""
|
|
if not window:
|
|
return None
|
|
toks = [w for w in re.sub(r"[^a-z0-9 ]", " ", (name or "").lower()).split()
|
|
if len(w) >= 4]
|
|
low = window.lower()
|
|
pos = -1
|
|
for tok in sorted(toks, key=len, reverse=True):
|
|
p = low.find(tok)
|
|
if p != -1:
|
|
pos = p; break
|
|
if pos == -1:
|
|
return window[:radius * 2].strip()
|
|
a = max(0, pos - radius)
|
|
return window[a:pos + radius].strip()
|
|
|
|
|
|
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
|
|
d_ev = [None] * len(distinct)
|
|
try:
|
|
if args.window and len(text) > args.window:
|
|
d_sup, n_win, n_fail, d_ev = 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
|
|
d_ev = [text if s else None for s in d_sup]
|
|
except Exception as e:
|
|
print(f" [{r.get('sample_id')}] ERROR: {e}")
|
|
d_sup = [False] * len(distinct); n_win = 0; n_fail = 1
|
|
# per distinct claim, cut a tight excerpt (sentence around the name) from its
|
|
# supporting window, so the training sample carries the evidence text, not the
|
|
# whole 120KB window.
|
|
d_exc = [_excerpt(d_ev[j], distinct[j]["o"]) if d_sup[j] else None
|
|
for j in range(len(distinct))]
|
|
ann = []
|
|
for i, t in enumerate(gold):
|
|
j = key_of[i]
|
|
ann.append({**t, "llm_grounded": bool(d_sup[j]), "evidence": d_exc[j]})
|
|
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
|
|
import os
|
|
rows = [json.loads(l) for l in open(args.inp, encoding="utf-8")]
|
|
if args.limit:
|
|
rows = rows[:args.limit]
|
|
# RESUME: skip samples already in the output (append mode). A server blip that
|
|
# kills the run no longer wastes completed work — just re-launch the same cmd.
|
|
done_ids = set()
|
|
if os.path.exists(args.out):
|
|
for l in open(args.out, encoding="utf-8"):
|
|
try:
|
|
done_ids.add(json.loads(l).get("sample_id"))
|
|
except Exception:
|
|
pass
|
|
todo = [r for r in rows if r.get("sample_id") not in done_ids]
|
|
print(f"resume: {len(done_ids)} done, {len(todo)} to do (of {len(rows)})")
|
|
done = tot_g = tot_t = 0
|
|
t0 = time.time()
|
|
workers = max(1, args.workers)
|
|
with open(args.out, "a", encoding="utf-8") as out, \
|
|
ThreadPoolExecutor(max_workers=workers) as ex:
|
|
futs = {ex.submit(_match_one, args, r): r for r in todo}
|
|
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(todo)}] {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} new 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, NO_THINK
|
|
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("--think", action="store_true",
|
|
help="keep reasoning thinking on (vllm); default off for speed")
|
|
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
|
|
NO_THINK = not args.think
|
|
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()
|