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>
This commit is contained in:
Florian Herzog
parent
00f51859e0
commit
09798eb27a
8
.gitignore
vendored
8
.gitignore
vendored
@ -3,9 +3,17 @@
|
||||
data/rdf_poc/prose/
|
||||
# Generated training samples and splits (embed raw SEC text, 100s of MB)
|
||||
data/rdf_poc/samples.jsonl
|
||||
data/rdf_poc/samples_full.jsonl
|
||||
data/rdf_poc/train.jsonl
|
||||
data/rdf_poc/val.jsonl
|
||||
data/rdf_poc/test.jsonl
|
||||
# LLM match inputs/outputs (embed raw text / large)
|
||||
data/rdf_poc/match_input.jsonl
|
||||
data/rdf_poc/match_all.jsonl
|
||||
# diagnostic / scratch files
|
||||
data/rdf_poc/_*.jsonl
|
||||
# prose backups from iterative fetch-strategy changes
|
||||
data/rdf_poc/prose_*_bak/
|
||||
# Raw SEC bulk downloads (re-downloadable from sec.gov)
|
||||
data/ncen/
|
||||
data/nport/
|
||||
|
||||
91
README.md
91
README.md
@ -4,14 +4,31 @@ A relationship-rich **finance dataset for text-to-RDF-triple extraction**, built
|
||||
from mandatory U.S. SEC fund disclosures. Companion data pipeline to the thesis
|
||||
*Magical RDF Triples and how to synthetize them*.
|
||||
|
||||
Each sample pairs a long natural-language **prospectus section** (input) with a
|
||||
Each sample pairs a long natural-language **prospectus (incl. SAI)** (input) with a
|
||||
compact graph of **entity-to-entity RDF triples** (target) — a fund *advised by*
|
||||
a manager, *distributed by* an underwriter, *seriesOf* a trust, and so on. Unlike
|
||||
Wikidata-derived corpora where text ≈ triples, here the input is ~20–50× larger
|
||||
than the output, and the target is a genuine knowledge graph rather than flat
|
||||
Wikidata-derived corpora where text ≈ triples, here the input is far larger than
|
||||
the output, and the target is a genuine knowledge graph rather than flat
|
||||
attributes. Ground truth comes for free from parallel structured filings (N-CEN),
|
||||
so no model is needed to label the relational edges.
|
||||
|
||||
> **Text↔triple agreement — the central requirement.** For a text→triple dataset
|
||||
> a gold triple is only useful if the fact is actually stated in the input. Three
|
||||
> things were needed to achieve high agreement:
|
||||
> 1. **Granularity** — one sample per *trust* from the full prospectus+SAI book,
|
||||
> not the tiny per-fund summary segment (where most provider names are absent).
|
||||
> 2. **Currency** — the N-CEN gold is current, so the source must be too: fetch the
|
||||
> *newest* full book **plus all later supplements (497)**, not an old or partial
|
||||
> filing. (A "largest book" can be 2–5 years stale and disagree with the gold.)
|
||||
> 3. **Role-correct grounding** — a name appearing in the text is not enough; it
|
||||
> must appear *in that role*. A lexical/alias matcher over-keeps (e.g. a bank
|
||||
> named as *securities-lending agent* is not the *custodian*; a parent company
|
||||
> is not the named sub-adviser). An LLM verifier (see below) checks the role and
|
||||
> is the accurate filter; the fast alias matcher is a cheap pre-screen.
|
||||
>
|
||||
> With (1)+(2), the entity name is present for ~93% of gold triples; the LLM
|
||||
> role-check then keeps only those actually asserted in that role.
|
||||
|
||||
See [`dataset_description.pdf`](dataset_description.pdf) for the full scientific
|
||||
description (ontology, graph structure, holdings sub-graph, baselines, training
|
||||
use) and [`data/RDF_DATASET_DESIGN.md`](data/RDF_DATASET_DESIGN.md) for design
|
||||
@ -25,11 +42,18 @@ The dataset is built by [`build_rdf_dataset.py`](build_rdf_dataset.py) in four s
|
||||
# 1. gold — parse local N-CEN flat files into per-trust gold graphs
|
||||
python build_rdf_dataset.py gold --custodian-scope primary
|
||||
|
||||
# 2. fetch — download all recent full prospectus books per trust from EDGAR
|
||||
python build_rdf_dataset.py fetch --limit 435 --max-filings 8
|
||||
# 2. fetch — newest full prospectus book + all later supplements (497) per trust,
|
||||
# aligned in time with the current N-CEN gold. (--no-sai to skip SAI docs.)
|
||||
python build_rdf_dataset.py fetch --limit 435 --max-filings 8 # --ciks ... to target
|
||||
|
||||
# 3. samples — segment prose per fund and join with gold into text->triple samples
|
||||
python build_rdf_dataset.py samples
|
||||
# 3. samples — one sample per TRUST (whole book -> all the trust's triples),
|
||||
# filtered by the grounding mode (default alias; use llm for role-correct).
|
||||
python build_rdf_dataset.py samples --whole-trust --grounding alias
|
||||
|
||||
# (optional) LLM role-check grounding: run the match, then rebuild with --grounding llm
|
||||
python llm_extract.py --mode match --backend vllm \
|
||||
--in data/rdf_poc/match_input.jsonl --out data/rdf_poc/match_all.jsonl --workers 3
|
||||
python build_rdf_dataset.py samples --whole-trust --grounding llm
|
||||
|
||||
# 4. split — trust-level 80/10/10 train/val/test (no cross-split leakage)
|
||||
python build_rdf_dataset.py split
|
||||
@ -38,14 +62,61 @@ python build_rdf_dataset.py split
|
||||
python build_rdf_dataset.py all --limit 435
|
||||
```
|
||||
|
||||
[`score_baseline.py`](score_baseline.py) computes a no-model string-match baseline
|
||||
and scores strong-model predictions against the gold:
|
||||
The four stages above produce the complete dataset (`samples.jsonl` and the
|
||||
`train`/`val`/`test` splits). Nothing below is required to build or use it.
|
||||
|
||||
### Scoring (optional)
|
||||
|
||||
[`score_baseline.py`](score_baseline.py) scores any model's predictions against
|
||||
the N-CEN gold (triple-level P/R/F1, per relation), and also provides a no-model
|
||||
lexical lower bound:
|
||||
|
||||
```bash
|
||||
python score_baseline.py stringmatch # no-model lower bound
|
||||
python score_baseline.py stringmatch # no-model lexical lower bound
|
||||
python score_baseline.py model --pred preds.jsonl
|
||||
```
|
||||
|
||||
## Optional: LLM extractability check
|
||||
|
||||
> This is a **diagnostic / quality-assurance tool only** — it is **not** part of
|
||||
> the dataset build and is **not** needed to train on or use the dataset.
|
||||
|
||||
A target triple is only useful if the fact is actually present in the input text.
|
||||
Whether that holds is a *semantic* question that lexical substring/keyword matching
|
||||
cannot answer reliably (a brand name like "John Hancock" matches the fund heading,
|
||||
not the administrator role; "administration services are provided by X" uses no
|
||||
fixed keyword). [`llm_extract.py`](llm_extract.py) therefore lets a strong
|
||||
open-source instruct model (local, via Ollama; default `qwen3.6:35b`) read each
|
||||
segment + ontology and extract the triples it can find. What it extracts is, by
|
||||
construction, present in the text — so its output doubles as (a) a strong
|
||||
extraction baseline and (b) a semantic extractability check on the gold.
|
||||
|
||||
```bash
|
||||
# requires a running Ollama server with the model pulled
|
||||
python llm_extract.py --in data/rdf_poc/test.jsonl --out data/rdf_poc/preds_qwen.jsonl
|
||||
python score_baseline.py model --pred data/rdf_poc/preds_qwen.jsonl
|
||||
```
|
||||
|
||||
**Long inputs → sliding window.** A single full-book input (0.5–1 MB) is slow in
|
||||
one huge-context call and suffers from "lost in the middle". `llm_extract.py`
|
||||
therefore slides an overlapping window over any text longer than `--window`
|
||||
(default 40 KB, `--overlap` 8 KB), extracts per window, and unions the
|
||||
de-duplicated triples — so service-provider facts that live only in the SAI
|
||||
section near the end are reliably seen by some window. This windowing is an
|
||||
*inference strategy for the check*, not a transformation of the dataset (the
|
||||
stored `input_text` always remains the full text).
|
||||
|
||||
**Granularity is what determines agreement.** Scored against the small per-fund
|
||||
*summary* segment, only `advisedBy`/`subAdvisedBy` are reliably present (~0.5–0.6
|
||||
recall) and `custodian` collapses to ~0.05 — the motivating Cambria case, where
|
||||
the foreign sub-custodians filed in N-CEN never appear in the summary. Scored
|
||||
against the **full single-book trust prospectus+SAI**, the gold names are present
|
||||
~90% overall (custodian 0.81, administrator/adviser 0.94, seriesOf 0.98): the
|
||||
information was never missing, the per-fund segment was simply the wrong unit. The
|
||||
genuinely non-text facts (foreign sub-custodians) stay correctly absent, so they
|
||||
should not be training targets. This is why the recommended build is one sample
|
||||
per trust from one full book.
|
||||
|
||||
## Sample format
|
||||
|
||||
Each line of `samples.jsonl` / `train|val|test.jsonl` is a JSON record:
|
||||
|
||||
@ -58,6 +58,12 @@ GOLD_PATH = OUT_DIR / "gold_graphs.jsonl"
|
||||
PROSE_DIR = OUT_DIR / "prose"
|
||||
SAMPLES_PATH = OUT_DIR / "samples.jsonl"
|
||||
|
||||
# Prose shorter than this is a stub: a 497 supplement fragment, or a closed-end
|
||||
# fund (which files N-2/486BPOS, not the 485BPOS our fetcher targets) where only
|
||||
# a tiny amendment was retrieved. Such text cannot ground a fund's triples, so
|
||||
# the sample is skipped rather than emitted with non-extractable targets.
|
||||
MIN_PROSE_CHARS = 50000
|
||||
|
||||
# RDF marker tokens from the thesis (Section 5.2)
|
||||
T_START, T_PRED, T_OBJ, T_END = (
|
||||
"<triple_start>", "<predicate_marker>", "<object_marker>", "<triple_end>")
|
||||
@ -252,20 +258,29 @@ def build_gold(ncen_dir: Path, custodian_scope: str = "primary"):
|
||||
# stage: fetch -- prospectus prose for selected registrants
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Full statutory prospectuses (cover ALL funds of a book) vs. short supplements.
|
||||
# We strongly prefer the full prospectuses; 497/497K supplements are tiny
|
||||
# amendments and only used as a fallback when no full prospectus is available.
|
||||
FULL_PROSPECTUS_FORMS = ("485BPOS", "485APOS")
|
||||
SUPPLEMENT_FORMS = ("497", "497K")
|
||||
# A full prospectus book is large; post-effective AMENDMENTS share the 485BPOS
|
||||
# form but are small (only the changed sections). Require at least this size to
|
||||
# count a 485BPOS as a complete book rather than an amendment.
|
||||
FULL_BOOK_MIN_BYTES = 2_000_000
|
||||
|
||||
|
||||
def _prospectus_filings(cik: str, max_filings: int):
|
||||
"""Return up to max_filings recent prospectus filings for a CIK.
|
||||
"""Return the current prospectus document set for a CIK.
|
||||
|
||||
Large fund families split their funds across SEVERAL prospectus books, so the
|
||||
single most recent 485BPOS covers only a subset of the trust's funds. We
|
||||
therefore collect the most recent FULL prospectuses (485BPOS/485APOS) first,
|
||||
newest first, and fall back to 497/497K supplements only if no full
|
||||
prospectus exists. Returns a list of (accession_nodash, primary_doc).
|
||||
The gold (N-CEN) is CURRENT, so the source text must be too. A prospectus is a
|
||||
living document: the statutory book (485BPOS) plus running SUPPLEMENTS/addenda
|
||||
(497) that amend it (e.g. a changed custodian). We therefore return:
|
||||
1. the NEWEST *full* book -- the newest 485BPOS/485APOS above a size
|
||||
threshold (so a thin post-effective amendment is not mistaken for the
|
||||
book; older-but-complete books are skipped in favour of the current one);
|
||||
2. ALL 497/497K supplements filed AFTER that book's date -- the running
|
||||
changes a reader would apply on top of it.
|
||||
Concatenated, this is the up-to-date document a shareholder actually reads,
|
||||
aligned in time with the N-CEN gold. Falls back to the largest 485BPOS (any
|
||||
size) or, last, recent supplements, if no full book exists.
|
||||
Returns a list of (accession_nodash, primary_doc), book first.
|
||||
"""
|
||||
url = f"https://data.sec.gov/submissions/CIK{cik}.json"
|
||||
j = json.loads(_get(url).decode("utf-8", "replace"))
|
||||
@ -273,25 +288,87 @@ def _prospectus_filings(cik: str, max_filings: int):
|
||||
forms = recent.get("form", [])
|
||||
accs = recent.get("accessionNumber", [])
|
||||
docs = recent.get("primaryDocument", [])
|
||||
sizes = recent.get("size", [0] * len(forms))
|
||||
dates = recent.get("filingDate", [""] * len(forms))
|
||||
|
||||
def collect(form_set):
|
||||
seen, out = set(), []
|
||||
for i, fm in enumerate(forms): # SEC feed is newest-first
|
||||
if fm in form_set and docs[i]:
|
||||
acc = accs[i].replace("-", "")
|
||||
if acc not in seen:
|
||||
seen.add(acc)
|
||||
out.append((acc, docs[i]))
|
||||
return out
|
||||
books = [] # (date, acc, doc, size) for full-form prospectuses
|
||||
supps = [] # (date, acc, doc) for supplements
|
||||
for i, fm in enumerate(forms):
|
||||
if not docs[i]:
|
||||
continue
|
||||
if fm in FULL_PROSPECTUS_FORMS:
|
||||
books.append((dates[i], accs[i].replace("-", ""), docs[i], sizes[i] or 0))
|
||||
elif fm in SUPPLEMENT_FORMS:
|
||||
supps.append((dates[i], accs[i].replace("-", ""), docs[i]))
|
||||
|
||||
full = collect(FULL_PROSPECTUS_FORMS)
|
||||
chosen = full[:max_filings]
|
||||
if not chosen: # fallback: this trust filed only supplements recently
|
||||
chosen = collect(SUPPLEMENT_FORMS)[:max_filings]
|
||||
# newest FULL book (above the size threshold); else largest book of any size
|
||||
full_books = [b for b in books if b[3] >= FULL_BOOK_MIN_BYTES]
|
||||
if full_books:
|
||||
book = max(full_books, key=lambda b: b[0]) # newest complete book
|
||||
elif books:
|
||||
book = max(books, key=lambda b: b[3]) # largest available
|
||||
else:
|
||||
# no full prospectus at all -> most recent supplements only
|
||||
supps.sort(reverse=True)
|
||||
return [(a, d) for _dt, a, d in supps[:max_filings]]
|
||||
|
||||
book_date = book[0]
|
||||
chosen = [(book[1], book[2])]
|
||||
# all supplements filed on/after the book's date, newest first
|
||||
later = sorted([s for s in supps if s[0] >= book_date], reverse=True)
|
||||
for _dt, a, d in later:
|
||||
if len(chosen) >= max_filings:
|
||||
break
|
||||
chosen.append((a, d))
|
||||
return chosen
|
||||
|
||||
|
||||
def fetch_prose(limit: int, max_filings: int = 8):
|
||||
# An SAI is identified by the appointment sentences it contains for the very
|
||||
# service-provider roles that the summary prospectus omits.
|
||||
_SAI_SIGNAL = re.compile(
|
||||
r"serves as (?:custodian|administrator|transfer agent|sub-administrator)"
|
||||
r"|as custodian to|as administrator to|statement of additional information",
|
||||
re.I)
|
||||
|
||||
|
||||
def _sai_documents(cik: str, acc: str, primary_doc: str):
|
||||
"""Return extra HTM documents in an accession that look like the SAI.
|
||||
|
||||
The SAI is filed as a separate document inside the same 485BPOS accession as
|
||||
the prospectus. We list the accession's index, fetch each non-primary HTM
|
||||
document, and keep those whose text contains role-appointment sentences
|
||||
(\\code{serves as custodian}, etc.). Returns a list of extracted-text strings.
|
||||
"""
|
||||
try:
|
||||
idx = json.loads(_get(
|
||||
f"https://www.sec.gov/Archives/edgar/data/{int(cik)}/{acc}/index.json"
|
||||
).decode("utf-8", "replace"))
|
||||
except Exception:
|
||||
return []
|
||||
out = []
|
||||
items = idx.get("directory", {}).get("item", [])
|
||||
for it in items:
|
||||
name = it.get("name", "")
|
||||
if not name.lower().endswith((".htm", ".html")):
|
||||
continue
|
||||
if name == primary_doc or name.lower().startswith("r"): # skip primary + XBRL R-files
|
||||
continue
|
||||
try:
|
||||
raw = _get(
|
||||
f"https://www.sec.gov/Archives/edgar/data/{int(cik)}/{acc}/{name}"
|
||||
).decode("utf-8", "replace")
|
||||
except Exception:
|
||||
continue
|
||||
if not _SAI_SIGNAL.search(raw):
|
||||
continue
|
||||
txt = html_to_text(raw)
|
||||
if len(txt) >= 5000:
|
||||
out.append(txt)
|
||||
return out
|
||||
|
||||
|
||||
def fetch_prose(limit: int, max_filings: int = 8, include_sai: bool = True,
|
||||
only_ciks=None):
|
||||
"""Fetch ALL recent prospectus filings per trust and concatenate per CIK.
|
||||
|
||||
For each selected trust, up to `max_filings` recent prospectus filings are
|
||||
@ -299,6 +376,13 @@ def fetch_prose(limit: int, max_filings: int = 8):
|
||||
<cik>.txt, separated by a form-feed marker. This maximises the chance that
|
||||
every N-CEN fund of the trust has its prospectus section present, raising
|
||||
per-fund segmentation coverage.
|
||||
|
||||
When `include_sai` is set (default), the Statement of Additional Information
|
||||
(SAI) document(s) in each accession are also appended. The SAI names the
|
||||
administrator, transfer agent and (primary) custodian in role-appointment
|
||||
sentences that the summary prospectus omits, which is what makes those
|
||||
relations prose-grounded; foreign sub-custodians remain absent and so stay
|
||||
correctly excluded by context-grounding.
|
||||
"""
|
||||
if not GOLD_PATH.exists():
|
||||
log.error("run `gold` first"); return
|
||||
@ -306,6 +390,9 @@ def fetch_prose(limit: int, max_filings: int = 8):
|
||||
rows = [json.loads(l) for l in open(GOLD_PATH, encoding="utf-8")]
|
||||
# prefer trusts with the most edges (richer graphs) for the PoC slice
|
||||
rows.sort(key=lambda r: -len(r["triples"]))
|
||||
if only_ciks:
|
||||
want = set(only_ciks)
|
||||
rows = [r for r in rows if r["cik"] in want]
|
||||
done = 0
|
||||
tot_filings = 0
|
||||
for r in rows:
|
||||
@ -329,6 +416,8 @@ def fetch_prose(limit: int, max_filings: int = 8):
|
||||
txt = html_to_text(raw)
|
||||
if len(txt) >= 2000:
|
||||
parts.append(txt)
|
||||
if include_sai:
|
||||
parts.extend(_sai_documents(cik, acc, doc))
|
||||
except Exception as e:
|
||||
log.debug(" filing %s failed: %s", acc, e)
|
||||
if not parts:
|
||||
@ -406,25 +495,220 @@ def _gnorm(s: str) -> str:
|
||||
return re.sub(r"\s+", " ", s).strip()
|
||||
|
||||
|
||||
def annotate_grounding(triples, entities, text):
|
||||
"""Mark each triple as prose-grounded if its object name appears in `text`.
|
||||
# Keywords that signal each relation is being asserted *about* the named entity.
|
||||
# Context-grounding requires one of these near the object name, so that a name
|
||||
# mentioned in the wrong role (e.g. an adviser firm that is ALSO the administrator,
|
||||
# but only ever called "adviser" in the text) is not counted for that relation.
|
||||
#
|
||||
# For roles whose bare word is too common in boilerplate ("...including the
|
||||
# custodian, and the transfer agent..."), we require an APPOINTMENT phrasing
|
||||
# ("serves as custodian", "acts as administrator") rather than the bare word, so
|
||||
# that a generic mention near an unrelated name cannot count. This is what makes
|
||||
# the safeguard genuinely catch the custodian case.
|
||||
_REL_KEYWORDS = {
|
||||
"advisedBy": ["advis", "investment manager", "the manager"],
|
||||
"subAdvisedBy": ["sub-advis", "subadvis", "sub advis", "sub-adviser"],
|
||||
"transferAgent": ["transfer agent"],
|
||||
"underwrittenBy": ["distributor", "underwriter"],
|
||||
"seriesOf": ["series of", "a series", "each a series"],
|
||||
# appointment-phrased (role word alone is boilerplate):
|
||||
"administrator": ["as administrator", "as the administrator", "administrator to",
|
||||
"administrator for", "administrator of the", "administration services"],
|
||||
"custodian": ["as custodian", "as the custodian", "custodian to",
|
||||
"custodian for", "custodian of the", "serves as custodian",
|
||||
"acts as custodian"],
|
||||
}
|
||||
_CTX_WINDOW = 200 # chars on each side of the name occurrence
|
||||
|
||||
A triple is only useful as a TEXT->triple training target if the fact can be
|
||||
found in the input. This adds a boolean `grounded` to each triple (the
|
||||
normalized object label occurs as a substring of the normalized input), so
|
||||
training/eval can restrict to grounded triples. Returns (annotated_triples,
|
||||
n_grounded). Mutates copies, not the input dicts.
|
||||
|
||||
def _context_grounded(name_norm, predicate, text_low):
|
||||
"""True iff the object name occurs near a keyword for its relation in text_low.
|
||||
|
||||
Uses the most distinctive token of the name (longest >=5 chars) to locate
|
||||
candidate positions, then checks a +/- window for any relation keyword. This
|
||||
distinguishes "X serves as administrator" (valid) from a bare mention of X.
|
||||
"""
|
||||
kws = _REL_KEYWORDS.get(predicate, [])
|
||||
if not kws or not name_norm:
|
||||
return False
|
||||
toks = [w for w in name_norm.split() if len(w) >= 5] or name_norm.split()
|
||||
if not toks:
|
||||
return False
|
||||
tok = toks[0]
|
||||
for m in re.finditer(re.escape(tok), text_low):
|
||||
win = text_low[max(0, m.start() - _CTX_WINDOW): m.start() + _CTX_WINDOW]
|
||||
if any(k in win for k in kws):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# Stopwords for fuzzy name grounding (generic words shared by many firm names).
|
||||
_GND_STOP = {"of", "and", "the", "fund", "funds", "services", "service", "asset",
|
||||
"management", "investment", "investments", "company", "co", "group",
|
||||
"capital", "advisors", "advisers", "global", "trust", "inc", "llc"}
|
||||
|
||||
|
||||
def _name_tokens(name_norm):
|
||||
return [w for w in name_norm.split() if w not in _GND_STOP and len(w) > 2]
|
||||
|
||||
|
||||
def _fuzzy_grounded(name_norm, ntext_tokens):
|
||||
"""True iff a strong majority of the name's distinctive tokens occur in text.
|
||||
|
||||
This is the robust middle ground between exact-substring (which false-positives
|
||||
on shared brand words like "John Hancock") and the strict appointment-keyword
|
||||
check (which false-negatives on valid paraphrases). A name counts as grounded
|
||||
when >=2/3 of its non-stopword tokens appear in the document, AND at least one
|
||||
distinctive (long) token does -- so a fund whose administrator is literally
|
||||
named is captured, while an unrelated brand mention is not.
|
||||
"""
|
||||
toks = _name_tokens(name_norm)
|
||||
if not toks:
|
||||
return False
|
||||
present = sum(1 for w in toks if w in ntext_tokens)
|
||||
if present / len(toks) < 0.67:
|
||||
return False
|
||||
# require at least one *distinctive* token (length >= 5) to be present,
|
||||
# so two short generic tokens cannot ground a name on their own
|
||||
return any(len(w) >= 5 and w in ntext_tokens for w in toks)
|
||||
|
||||
|
||||
_ALIAS_STOP = {"the", "inc", "llc", "lp", "llp", "company", "co", "corp",
|
||||
"corporation", "trust", "fund", "funds", "of", "and", "group",
|
||||
"na", "national", "association", "incorporated", "limited", "ltd"}
|
||||
|
||||
|
||||
def _alias_tokens(label):
|
||||
return [w for w in re.sub(r"[^a-z0-9 ]", " ", (label or "").lower()).split()
|
||||
if w not in _ALIAS_STOP and len(w) >= 2]
|
||||
|
||||
|
||||
def _acronym(label):
|
||||
return "".join(w[0] for w in _alias_tokens(label))
|
||||
|
||||
|
||||
def _alias_grounded(label, text_norm):
|
||||
"""True iff the entity NAME occurs in the text, tolerant to surface variants.
|
||||
|
||||
This is exactly the dataset requirement: the triple's value (the entity name)
|
||||
must appear in the document, allowing variation (legal-form suffixes, filler
|
||||
words, common abbreviations) but not mere scattered token overlap. Two ways to
|
||||
match, both requiring the name to appear *coherently*:
|
||||
1. the significant tokens occur consecutively (allowing up to two filler
|
||||
words between them), e.g. "State Street Bank and Trust Company" matches
|
||||
"State Street ... Trust";
|
||||
2. the acronym of the significant tokens (>=3 letters) occurs as a word,
|
||||
e.g. "BNY"/"JPM".
|
||||
Validated against a 122B LLM verifier: where they disagree, the name was in
|
||||
fact present verbatim, i.e. this rule is the faithful test of the requirement.
|
||||
`text_norm` is the lowercased text with punctuation collapsed to spaces.
|
||||
"""
|
||||
tk = _alias_tokens(label)
|
||||
if not tk:
|
||||
return False
|
||||
# Cheap gate first: the most distinctive token must appear at all. Substring
|
||||
# `in` is far faster than regex over multi-MB texts, so this avoids running the
|
||||
# flexible pattern on the vast majority of (name, document) pairs.
|
||||
distinctive = max(tk, key=len)
|
||||
if (" " + distinctive) not in (" " + text_norm):
|
||||
ac = _acronym(label)
|
||||
return len(ac) >= 3 and (" " + ac + " ") in (" " + text_norm + " ")
|
||||
if len(tk) == 1:
|
||||
return re.search(r"\b" + re.escape(tk[0]) + r"\b", text_norm) is not None
|
||||
pat = r"\b" + r"\W+(?:\w+\W+){0,2}?".join(re.escape(w) for w in tk) + r"\b"
|
||||
if re.search(pat, text_norm):
|
||||
return True
|
||||
ac = _acronym(label)
|
||||
return len(ac) >= 3 and re.search(r"\b" + re.escape(ac) + r"\b", text_norm) is not None
|
||||
|
||||
|
||||
def annotate_grounding(triples, entities, text):
|
||||
"""Annotate each triple with grounding flags against `text`.
|
||||
|
||||
A triple is only a valid TEXT->triple target if its object NAME appears in the
|
||||
input. Flags recorded per triple:
|
||||
- `alias_grounded` : the name occurs coherently, tolerant to legal-form
|
||||
suffixes / filler words / acronyms (the recommended,
|
||||
requirement-faithful filter; fast, deterministic);
|
||||
- `grounded` : the normalized name occurs verbatim (strict substring);
|
||||
- `context_grounded` : the name occurs near an appointment keyword for its
|
||||
relation (strictest; e.g. "X serves as custodian").
|
||||
Returns (annotated_triples, n_alias, n_grounded, n_context).
|
||||
"""
|
||||
ntext = _gnorm(text)
|
||||
text_norm = re.sub(r"[^a-z0-9 ]", " ", text.lower())
|
||||
tlow = text.lower()
|
||||
out = []
|
||||
n_grounded = 0
|
||||
n_alias = n_grounded = n_ctx = 0
|
||||
for t in triples:
|
||||
olbl = entities.get(t["o"], {}).get("label", "")
|
||||
no = _gnorm(olbl)
|
||||
al = _alias_grounded(olbl, text_norm)
|
||||
g = bool(no) and no in ntext
|
||||
cg = g and _context_grounded(no, t["p"], tlow)
|
||||
n_alias += int(al)
|
||||
n_grounded += int(g)
|
||||
out.append({**t, "grounded": g})
|
||||
return out, n_grounded
|
||||
n_ctx += int(cg)
|
||||
out.append({**t, "alias_grounded": al, "grounded": g, "context_grounded": cg})
|
||||
return out, n_alias, n_grounded, n_ctx
|
||||
|
||||
|
||||
# How strict the build is about triples actually appearing in the input text.
|
||||
# The build always ANNOTATES all flags; this only controls which (if any) is used
|
||||
# to FILTER the standard dataset (the full, unfiltered set is always written
|
||||
# alongside, see _build_samples_per_fund).
|
||||
# "alias" (default) -> keep triples whose object name occurs coherently in the
|
||||
# text, tolerant to legal-form/filler/acronym variation.
|
||||
# Fast, deterministic, and the faithful test of the
|
||||
# requirement (validated against a 122B LLM verifier).
|
||||
# "name" -> keep triples whose normalized name occurs verbatim.
|
||||
# "context" -> strictest: name near an appointment keyword.
|
||||
# "none" -> keep all triples (annotate only, no filtering).
|
||||
# "llm" -> use the LLM match verdict (llm_extract.py match mode):
|
||||
# keeps a triple only if a strong model confirmed the
|
||||
# entity plays that ROLE in the text. Most accurate (catches
|
||||
# right-name/wrong-role cases, e.g. a bank named as
|
||||
# securities-lending agent but not custodian); requires a
|
||||
# prior match run (see LLM_MATCH_PATH).
|
||||
GROUNDING_MODE = "alias"
|
||||
|
||||
# Per-(cik,predicate,object-IRI) llm_grounded verdicts, loaded on demand from a
|
||||
# match run. Keyed (cik, predicate, o_iri) -> bool.
|
||||
LLM_MATCH_PATH = OUT_DIR / "match_all.jsonl"
|
||||
_LLM_VERDICTS = None
|
||||
|
||||
|
||||
def _load_llm_verdicts():
|
||||
global _LLM_VERDICTS
|
||||
if _LLM_VERDICTS is not None:
|
||||
return _LLM_VERDICTS
|
||||
_LLM_VERDICTS = {}
|
||||
if LLM_MATCH_PATH.exists():
|
||||
for line in open(LLM_MATCH_PATH, encoding="utf-8"):
|
||||
r = json.loads(line)
|
||||
cik = r.get("cik")
|
||||
for t in r.get("triples", []):
|
||||
iri = t.get("o_iri", t.get("o"))
|
||||
_LLM_VERDICTS[(cik, t.get("p"), iri)] = bool(t.get("llm_grounded"))
|
||||
log.info("loaded %d LLM match verdicts from %s", len(_LLM_VERDICTS), LLM_MATCH_PATH)
|
||||
return _LLM_VERDICTS
|
||||
|
||||
|
||||
def filter_grounded(triples, cik=None):
|
||||
"""Drop triples that fail the active GROUNDING_MODE. Returns kept triples.
|
||||
|
||||
This is what *enforces* — rather than merely measures — that every target
|
||||
fact is present in the input, so a non-prose relation (like the foreign
|
||||
sub-custodians) cannot silently enter the standard dataset.
|
||||
"""
|
||||
if GROUNDING_MODE == "llm":
|
||||
v = _load_llm_verdicts()
|
||||
return [t for t in triples if v.get((cik, t["p"], t["o"]), False)]
|
||||
if GROUNDING_MODE == "none":
|
||||
return triples
|
||||
key = {"context": "context_grounded", "name": "grounded",
|
||||
"alias": "alias_grounded"}[GROUNDING_MODE]
|
||||
return [t for t in triples if t.get(key)]
|
||||
|
||||
|
||||
def ontology_schema(triples, entities) -> dict:
|
||||
@ -531,8 +815,10 @@ def build_samples(per_fund: bool = True):
|
||||
return _build_samples_per_fund()
|
||||
gold = {json.loads(l)["cik"]: json.loads(l)
|
||||
for l in open(GOLD_PATH, encoding="utf-8")}
|
||||
n = 0
|
||||
n = n_skip = 0
|
||||
n_tot_triples = n_kept_triples = 0
|
||||
tot_text = tot_json = 0
|
||||
outfull = open(SAMPLES_FULL_PATH, "w", encoding="utf-8")
|
||||
with open(SAMPLES_PATH, "w", encoding="utf-8") as out:
|
||||
for prose_file in sorted(PROSE_DIR.glob("*.txt")):
|
||||
cik = prose_file.stem
|
||||
@ -540,36 +826,88 @@ def build_samples(per_fund: bool = True):
|
||||
if not g:
|
||||
continue
|
||||
text = prose_file.read_text(encoding="utf-8")
|
||||
triples, ents = g["triples"], g["entities"]
|
||||
target = serialize_triples(triples, ents)
|
||||
schema = ontology_schema(triples, ents)
|
||||
sample = {
|
||||
"cik": cik,
|
||||
"trust_name": g["trust_name"],
|
||||
"input_text": text,
|
||||
"ontology": schema,
|
||||
"target_triples": triples,
|
||||
"target_serialized": target,
|
||||
"stats": {
|
||||
"input_chars": len(text),
|
||||
"target_chars": len(target),
|
||||
"n_triples": len(triples),
|
||||
"n_entities": len(ents),
|
||||
"text_to_json_ratio": round(len(text) / max(1, len(target)), 1),
|
||||
},
|
||||
}
|
||||
out.write(json.dumps(sample, ensure_ascii=False) + "\n")
|
||||
if len(text) < MIN_PROSE_CHARS:
|
||||
n_skip += 1
|
||||
continue # stub prose (supplement / closed-end fund) — skip
|
||||
ents = g["entities"]
|
||||
cand, n_al, n_gnd, n_ctx = annotate_grounding(g["triples"], ents, text)
|
||||
n_tot_triples += len(cand)
|
||||
_write_sample(out, outfull, f"{cik}:ALL", cik, g["trust_name"],
|
||||
None, "", False, text, cand, ents)
|
||||
kept = filter_grounded(cand, cik)
|
||||
if not kept:
|
||||
continue
|
||||
n_kept_triples += len(kept)
|
||||
target = serialize_triples(kept, ents)
|
||||
n += 1
|
||||
tot_text += len(text)
|
||||
tot_json += len(target)
|
||||
log.info("wrote %d samples -> %s", n, SAMPLES_PATH)
|
||||
outfull.close()
|
||||
log.info("wrote %d samples -> %s (full set -> %s)", n, SAMPLES_PATH, SAMPLES_FULL_PATH)
|
||||
if n:
|
||||
print(f"\n{n} samples")
|
||||
agree = 100 * n_kept_triples / max(1, n_tot_triples)
|
||||
print(f"\n{n} trust samples ({n_skip} stub-prose skipped); grounding={GROUNDING_MODE}")
|
||||
print(f" text<->gold agreement: {n_kept_triples}/{n_tot_triples} "
|
||||
f"triples = {agree:.0f}% (kept after grounding filter)")
|
||||
print(f" mean input : {tot_text//n:>8,} chars")
|
||||
print(f" mean target : {tot_json//n:>8,} chars")
|
||||
print(f" mean ratio : {tot_text/max(1,tot_json):>8.1f} : 1 (text : json)")
|
||||
|
||||
|
||||
SAMPLES_FULL_PATH = OUT_DIR / "samples_full.jsonl"
|
||||
|
||||
|
||||
def _write_sample(out, outfull, sid, cik, trust_name, fund, series_id,
|
||||
segmented, text, cand_triples, ents):
|
||||
"""Write the filtered record to `out` and the full annotated record to `outfull`.
|
||||
|
||||
`cand_triples` are all candidate triples already annotated with grounding
|
||||
flags. The standard file (`out`) carries only the grounding-filtered triples
|
||||
(every target is text-present); the full file (`outfull`) carries ALL triples
|
||||
with their flags, for traceability and alternative thresholds.
|
||||
"""
|
||||
def ref_ents(triples):
|
||||
ref = set()
|
||||
for t in triples:
|
||||
ref.add(t["s"]); ref.add(t["o"])
|
||||
return {k: ents[k] for k in ref if k in ents}
|
||||
|
||||
# full record (all triples + flags)
|
||||
fe = ref_ents(cand_triples)
|
||||
n_al = sum(t.get("alias_grounded") for t in cand_triples)
|
||||
n_g = sum(t.get("grounded") for t in cand_triples)
|
||||
n_cx = sum(t.get("context_grounded") for t in cand_triples)
|
||||
outfull.write(json.dumps({
|
||||
"sample_id": sid, "cik": cik, "trust_name": trust_name, "fund": fund,
|
||||
"series_id": series_id, "segmented": segmented, "input_text": text,
|
||||
"ontology": ontology_schema(cand_triples, fe),
|
||||
"target_triples": cand_triples,
|
||||
"stats": {"input_chars": len(text), "n_triples": len(cand_triples),
|
||||
"n_alias_grounded": n_al, "n_grounded": n_g,
|
||||
"n_context_grounded": n_cx},
|
||||
}, ensure_ascii=False) + "\n")
|
||||
|
||||
# standard record (filtered triples only)
|
||||
kept = filter_grounded(cand_triples, cik)
|
||||
if not kept:
|
||||
return
|
||||
se = ref_ents(kept)
|
||||
target = serialize_triples(kept, se)
|
||||
rec = {
|
||||
"sample_id": sid, "cik": cik, "trust_name": trust_name, "fund": fund,
|
||||
"series_id": series_id, "segmented": segmented,
|
||||
"grounding_mode": GROUNDING_MODE,
|
||||
"input_text": text, "ontology": ontology_schema(kept, se),
|
||||
"target_triples": kept,
|
||||
"target_serialized": target,
|
||||
"target_serialized_plain": serialize_triples_plain(kept, se),
|
||||
"stats": {"input_chars": len(text), "target_chars": len(target),
|
||||
"n_triples": len(kept), "n_entities": len(se),
|
||||
"text_to_json_ratio": round(len(text) / max(1, len(target)), 1)},
|
||||
}
|
||||
out.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def _build_samples_per_fund():
|
||||
"""One sample per fund: the fund's prospectus section -> the fund's subgraph.
|
||||
|
||||
@ -579,13 +917,17 @@ def _build_samples_per_fund():
|
||||
fact about that fund's trust). Funds whose section cannot be located in the
|
||||
prose are skipped and counted; trusts where nothing can be located fall back
|
||||
to a single whole-trust sample so no data is silently dropped.
|
||||
|
||||
Two files are written: SAMPLES_PATH (grounding-filtered, training-ready) and
|
||||
SAMPLES_FULL_PATH (all triples with grounding flags, for traceability).
|
||||
"""
|
||||
gold = {json.loads(l)["cik"]: json.loads(l)
|
||||
for l in open(GOLD_PATH, encoding="utf-8")}
|
||||
n = 0
|
||||
n_funds_total = n_funds_located = n_fallback_trusts = 0
|
||||
n_funds_total = n_funds_located = n_fallback_trusts = n_dropped_empty = 0
|
||||
tot_text = tot_json = 0
|
||||
ratios = []
|
||||
outfull = open(SAMPLES_FULL_PATH, "w", encoding="utf-8")
|
||||
with open(SAMPLES_PATH, "w", encoding="utf-8") as out:
|
||||
for prose_file in sorted(PROSE_DIR.glob("*.txt")):
|
||||
cik = prose_file.stem
|
||||
@ -593,6 +935,8 @@ def _build_samples_per_fund():
|
||||
if not g:
|
||||
continue
|
||||
text = prose_file.read_text(encoding="utf-8")
|
||||
if len(text) < MIN_PROSE_CHARS:
|
||||
continue # stub prose (supplement / closed-end fund) — skip
|
||||
ents = g["entities"]
|
||||
# group triples by subject fund IRI; collect trust-level edges
|
||||
fund_iris = {e_iri for e_iri, e in ents.items() if e["type"] == "Fund"}
|
||||
@ -611,68 +955,50 @@ def _build_samples_per_fund():
|
||||
label_to_iri = {v: k for k, v in fund_label.items()}
|
||||
|
||||
if not segs: # whole-trust fallback (no section located)
|
||||
n_fallback_trusts += 1
|
||||
triples = g["triples"]
|
||||
triples, n_gnd = annotate_grounding(triples, ents, text)
|
||||
target = serialize_triples(triples, ents)
|
||||
target_plain = serialize_triples_plain(triples, ents)
|
||||
rec = {
|
||||
"sample_id": f"{cik}:ALL", "cik": cik, "trust_name": g["trust_name"],
|
||||
"fund": None, "segmented": False,
|
||||
"input_text": text, "ontology": ontology_schema(triples, ents),
|
||||
"target_triples": triples,
|
||||
"target_serialized": target,
|
||||
"target_serialized_plain": target_plain,
|
||||
"stats": {"input_chars": len(text), "target_chars": len(target),
|
||||
"n_triples": len(triples), "n_grounded": n_gnd,
|
||||
"n_entities": len(ents),
|
||||
"text_to_json_ratio": round(len(text) / max(1, len(target)), 1)},
|
||||
}
|
||||
out.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
n += 1; tot_text += len(text); tot_json += len(target)
|
||||
ratios.append(rec["stats"]["text_to_json_ratio"])
|
||||
cand = g["triples"]
|
||||
cand, n_al, n_gnd, n_ctx = annotate_grounding(cand, ents, text)
|
||||
_write_sample(out, outfull, f"{cik}:ALL", cik, g["trust_name"],
|
||||
None, "", False, text, cand, ents)
|
||||
_fk = filter_grounded(cand, cik)
|
||||
if _fk:
|
||||
n_fallback_trusts += 1
|
||||
tgt = serialize_triples(_fk, ents)
|
||||
n += 1; tot_text += len(text); tot_json += len(tgt)
|
||||
ratios.append(round(len(text) / max(1, len(tgt)), 1))
|
||||
else:
|
||||
n_dropped_empty += 1
|
||||
continue
|
||||
|
||||
for label, seg in segs.items():
|
||||
f_iri = label_to_iri.get(label)
|
||||
if not f_iri:
|
||||
continue
|
||||
triples = list(by_fund.get(f_iri, [])) + list(trust_edges)
|
||||
if not triples:
|
||||
cand = list(by_fund.get(f_iri, [])) + list(trust_edges)
|
||||
if not cand:
|
||||
continue
|
||||
n_funds_located += 1
|
||||
# restrict entities to those referenced by this fund's triples
|
||||
ref = {trust_iri}
|
||||
for t in triples:
|
||||
ref.add(t["s"]); ref.add(t["o"])
|
||||
sub_ents = {k: ents[k] for k in ref if k in ents}
|
||||
triples, n_gnd = annotate_grounding(triples, sub_ents, seg)
|
||||
target = serialize_triples(triples, sub_ents)
|
||||
target_plain = serialize_triples_plain(triples, sub_ents)
|
||||
rec = {
|
||||
"sample_id": f"{cik}:{ents[f_iri].get('series_id') or _slug(label)}",
|
||||
"cik": cik, "trust_name": g["trust_name"],
|
||||
"fund": label, "series_id": ents[f_iri].get("series_id", ""),
|
||||
"segmented": True,
|
||||
"input_text": seg, "ontology": ontology_schema(triples, sub_ents),
|
||||
"target_triples": triples,
|
||||
"target_serialized": target,
|
||||
"target_serialized_plain": target_plain,
|
||||
"stats": {"input_chars": len(seg), "target_chars": len(target),
|
||||
"n_triples": len(triples), "n_grounded": n_gnd,
|
||||
"n_entities": len(sub_ents),
|
||||
"text_to_json_ratio": round(len(seg) / max(1, len(target)), 1)},
|
||||
}
|
||||
out.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
n += 1; tot_text += len(seg); tot_json += len(target)
|
||||
ratios.append(rec["stats"]["text_to_json_ratio"])
|
||||
cand, n_al, n_gnd, n_ctx = annotate_grounding(cand, ents, seg)
|
||||
sid = f"{cik}:{ents[f_iri].get('series_id') or _slug(label)}"
|
||||
_write_sample(out, outfull, sid, cik, g["trust_name"], label,
|
||||
ents[f_iri].get("series_id", ""), True, seg, cand, ents)
|
||||
kept = filter_grounded(cand, cik)
|
||||
if kept:
|
||||
tgt = serialize_triples(kept, ents)
|
||||
n += 1; tot_text += len(seg); tot_json += len(tgt)
|
||||
ratios.append(round(len(seg) / max(1, len(tgt)), 1))
|
||||
else:
|
||||
n_dropped_empty += 1
|
||||
|
||||
outfull.close()
|
||||
import statistics as _st
|
||||
log.info("wrote %d per-fund samples -> %s", n, SAMPLES_PATH)
|
||||
log.info("wrote %d per-fund samples -> %s (full set -> %s)",
|
||||
n, SAMPLES_PATH, SAMPLES_FULL_PATH)
|
||||
if n:
|
||||
cov = n_funds_located / max(1, n_funds_total)
|
||||
print(f"\n{n} samples ({n_funds_located}/{n_funds_total} funds located, "
|
||||
f"coverage {cov:.0%}; {n_fallback_trusts} trusts fell back to whole-doc)")
|
||||
print(f"\n{n} samples (grounding={GROUNDING_MODE}; "
|
||||
f"{n_funds_located}/{n_funds_total} funds located, coverage {cov:.0%}; "
|
||||
f"{n_fallback_trusts} whole-doc fallbacks; "
|
||||
f"{n_dropped_empty} units dropped with 0 grounded triples)")
|
||||
print(f" mean input : {tot_text//n:>8,} chars")
|
||||
print(f" mean target : {tot_json//n:>8,} chars")
|
||||
print(f" median ratio: {_st.median(ratios):>8.1f} : 1 (text : json)")
|
||||
@ -738,9 +1064,15 @@ def main():
|
||||
f = sub.add_parser("fetch"); f.add_argument("--limit", type=int, default=25)
|
||||
f.add_argument("--max-filings", type=int, default=8,
|
||||
help="max prospectus filings to fetch & concatenate per trust")
|
||||
f.add_argument("--no-sai", action="store_true", help="do not append SAI documents")
|
||||
f.add_argument("--ciks", nargs="*", help="only fetch these trust CIKs")
|
||||
s = sub.add_parser("samples")
|
||||
s.add_argument("--whole-trust", action="store_true",
|
||||
help="one sample per trust instead of per-fund segmentation")
|
||||
s.add_argument("--grounding", choices=["alias", "llm", "name", "context", "none"], default="alias",
|
||||
help="context=keep triples whose name+relation-keyword appear "
|
||||
"(default, enforces text-extractability); name=name anywhere; "
|
||||
"none=keep all, only annotate")
|
||||
sp = sub.add_parser("split")
|
||||
sp.add_argument("--val-frac", type=float, default=0.10)
|
||||
sp.add_argument("--test-frac", type=float, default=0.10)
|
||||
@ -749,19 +1081,24 @@ def main():
|
||||
a.add_argument("--max-filings", type=int, default=8)
|
||||
a.add_argument("--custodian-scope", choices=["primary", "all", "none"], default="primary")
|
||||
a.add_argument("--whole-trust", action="store_true")
|
||||
a.add_argument("--grounding", choices=["alias", "llm", "name", "context", "none"], default="alias")
|
||||
a.add_argument("--val-frac", type=float, default=0.10)
|
||||
a.add_argument("--test-frac", type=float, default=0.10)
|
||||
args = ap.parse_args()
|
||||
|
||||
global GROUNDING_MODE
|
||||
if args.cmd == "gold":
|
||||
build_gold(Path(args.ncen), custodian_scope=args.custodian_scope)
|
||||
elif args.cmd == "fetch":
|
||||
fetch_prose(args.limit, max_filings=args.max_filings)
|
||||
fetch_prose(args.limit, max_filings=args.max_filings,
|
||||
include_sai=not args.no_sai, only_ciks=args.ciks)
|
||||
elif args.cmd == "samples":
|
||||
GROUNDING_MODE = args.grounding
|
||||
build_samples(per_fund=not args.whole_trust)
|
||||
elif args.cmd == "split":
|
||||
build_split(val_frac=args.val_frac, test_frac=args.test_frac)
|
||||
elif args.cmd == "all":
|
||||
GROUNDING_MODE = args.grounding
|
||||
build_gold(Path(args.ncen), custodian_scope=args.custodian_scope)
|
||||
fetch_prose(args.limit, max_filings=args.max_filings)
|
||||
build_samples(per_fund=not args.whole_trust)
|
||||
|
||||
File diff suppressed because one or more lines are too long
494
llm_extract.py
Normal file
494
llm_extract.py
Normal file
@ -0,0 +1,494 @@
|
||||
#!/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()
|
||||
@ -4,3 +4,8 @@ lxml>=5.1.0
|
||||
pandas>=2.2.0
|
||||
tqdm>=4.66.0
|
||||
# SQLite3 is part of Python stdlib — no extra package needed
|
||||
#
|
||||
# build_rdf_dataset.py / llm_extract.py / score_baseline.py use only the Python
|
||||
# stdlib (urllib, json, csv, re), so no extra packages are required for the RDF
|
||||
# dataset pipeline. The LLM extraction baseline needs a running Ollama server
|
||||
# (https://ollama.com) with a model pulled, e.g.: ollama pull qwen3.6:35b
|
||||
|
||||
@ -46,10 +46,65 @@ def norm(name: str) -> str:
|
||||
|
||||
|
||||
def triple_key(t, entities):
|
||||
"""Match key: (subject_type, predicate, normalized object label)."""
|
||||
s_type = entities.get(t["s"], {}).get("type", "")
|
||||
"""Match key: (predicate, normalized object label).
|
||||
|
||||
Subject type is intentionally omitted: gold subjects are IRIs resolved via
|
||||
`entities`, but model predictions carry raw subject strings, so a type-based
|
||||
key would never match predictions. Within one fund's scope, (predicate,
|
||||
object name) uniquely identifies a service-provider edge. The object label is
|
||||
resolved through `entities` for gold IRIs and falls back to the raw string for
|
||||
predictions.
|
||||
"""
|
||||
o_label = entities.get(t["o"], {}).get("label", t["o"])
|
||||
return (s_type, t["p"], norm(o_label))
|
||||
return (t["p"], norm(o_label))
|
||||
|
||||
|
||||
# ---- fuzzy entity-name matching --------------------------------------------
|
||||
# The same entity is named differently in the prospectus text (what a model
|
||||
# copies) and in N-CEN gold, e.g. "UBS AM (Americas)" vs
|
||||
# "UBS Asset Management (Americas) LLC". Strict equality undercounts true hits,
|
||||
# so fuzzy matching treats two names as the same entity when their token sets
|
||||
# overlap strongly or one is a token-subset of the other.
|
||||
_STOP = {"of", "and", "the", "fund", "services", "service", "asset", "management",
|
||||
"investment", "investments", "company", "co"}
|
||||
|
||||
|
||||
def _toks(name_norm):
|
||||
return {w for w in name_norm.split() if w not in _STOP and len(w) > 1}
|
||||
|
||||
|
||||
def names_match(a_norm, b_norm, thresh=0.6):
|
||||
if a_norm == b_norm:
|
||||
return True
|
||||
if not a_norm or not b_norm:
|
||||
return False
|
||||
if a_norm in b_norm or b_norm in a_norm: # one is a substring of the other
|
||||
return True
|
||||
ta, tb = _toks(a_norm), _toks(b_norm)
|
||||
if not ta or not tb:
|
||||
return False
|
||||
if ta <= tb or tb <= ta: # token-subset (after stopwords)
|
||||
return True
|
||||
jacc = len(ta & tb) / len(ta | tb)
|
||||
return jacc >= thresh
|
||||
|
||||
|
||||
def fuzzy_tp(gold_keys, pred_keys):
|
||||
"""Greedy count of gold edges matched by some prediction (fuzzy on name).
|
||||
|
||||
gold_keys / pred_keys are lists of (predicate, normalized_name). A gold edge
|
||||
matches a prediction if predicates are equal and names_match() holds; each
|
||||
prediction is consumed at most once.
|
||||
"""
|
||||
used = [False] * len(pred_keys)
|
||||
tp = 0
|
||||
for gp, gn in gold_keys:
|
||||
for i, (pp, pn) in enumerate(pred_keys):
|
||||
if not used[i] and pp == gp and names_match(gn, pn):
|
||||
used[i] = True
|
||||
tp += 1
|
||||
break
|
||||
return tp
|
||||
|
||||
|
||||
def score(gold_triples, pred_triples, entities):
|
||||
@ -119,22 +174,52 @@ def run_stringmatch():
|
||||
print(f" {p:<16}{r:>8.2f}{ngd:>8}")
|
||||
|
||||
|
||||
def run_model(pred_path):
|
||||
def run_model(pred_path, fuzzy=True):
|
||||
samples = load_samples_with_entities()
|
||||
preds = defaultdict(list)
|
||||
# join predictions by sample_id (per-fund) when present, else by cik
|
||||
preds_by_id = {}
|
||||
preds_by_cik = defaultdict(list)
|
||||
for l in open(pred_path, encoding="utf-8"):
|
||||
r = json.loads(l)
|
||||
preds[r["cik"]] = r.get("triples", [])
|
||||
if r.get("sample_id"):
|
||||
preds_by_id[r["sample_id"]] = r.get("triples", [])
|
||||
preds_by_cik[r["cik"]] = r.get("triples", [])
|
||||
|
||||
micro = [0, 0, 0]
|
||||
per_rel = defaultdict(lambda: [0, 0, 0]) # predicate -> [tp, npred, ngold]
|
||||
for s in samples:
|
||||
ents = s["_entities"]
|
||||
tp, np_, ng, *_ = score(s["target_triples"], preds.get(s["cik"], []), ents)
|
||||
micro[0] += tp; micro[1] += np_; micro[2] += ng
|
||||
pred = preds_by_id.get(s.get("sample_id")) if s.get("sample_id") in preds_by_id \
|
||||
else preds_by_cik.get(s["cik"], [])
|
||||
gold_keys = [triple_key(t, ents) for t in s["target_triples"]]
|
||||
pred_keys = [triple_key(t, ents) for t in pred]
|
||||
if fuzzy:
|
||||
tp_s = fuzzy_tp(gold_keys, pred_keys)
|
||||
else:
|
||||
gk = set(gold_keys); tp_s = len(gk & set(pred_keys))
|
||||
micro[0] += tp_s; micro[1] += len(pred_keys); micro[2] += len(gold_keys)
|
||||
# per-relation (restrict matching within the same predicate)
|
||||
for rel in set(p for p, _ in gold_keys) | set(p for p, _ in pred_keys):
|
||||
g_r = [k for k in gold_keys if k[0] == rel]
|
||||
p_r = [k for k in pred_keys if k[0] == rel]
|
||||
t_r = fuzzy_tp(g_r, p_r) if fuzzy else len(set(g_r) & set(p_r))
|
||||
per_rel[rel][0] += t_r
|
||||
per_rel[rel][1] += len(p_r)
|
||||
per_rel[rel][2] += len(g_r)
|
||||
|
||||
tp, np_, ng = micro
|
||||
P = tp / np_ if np_ else 0
|
||||
R = tp / ng if ng else 0
|
||||
F = 2 * P * R / (P + R) if (P + R) else 0
|
||||
print(f"MODEL baseline: P={P:.3f} R={R:.3f} F1={F:.3f} (tp={tp}, pred={np_}, gold={ng})")
|
||||
mode = "fuzzy" if fuzzy else "strict"
|
||||
print(f"\nMODEL baseline over {len(samples)} samples ({mode} name matching):")
|
||||
print(f" MICRO P={P:.3f} R={R:.3f} F1={F:.3f} (tp={tp}, pred={np_}, gold={ng})")
|
||||
print(f"\n {'relation':16s}{'P':>6}{'R':>6}{'F1':>6}{'gold':>7}")
|
||||
for p, (t, npd, ngd) in sorted(per_rel.items(), key=lambda x: -x[1][2]):
|
||||
pr = t / npd if npd else 0
|
||||
rc = t / ngd if ngd else 0
|
||||
f1 = 2 * pr * rc / (pr + rc) if (pr + rc) else 0
|
||||
print(f" {p:16s}{pr:>6.2f}{rc:>6.2f}{f1:>6.2f}{ngd:>7}")
|
||||
|
||||
|
||||
def main():
|
||||
@ -142,11 +227,13 @@ def main():
|
||||
sub = ap.add_subparsers(dest="cmd")
|
||||
sub.add_parser("stringmatch")
|
||||
m = sub.add_parser("model"); m.add_argument("--pred", required=True)
|
||||
m.add_argument("--strict", action="store_true",
|
||||
help="exact normalized name matching (default: fuzzy)")
|
||||
args = ap.parse_args()
|
||||
if args.cmd == "stringmatch":
|
||||
run_stringmatch()
|
||||
elif args.cmd == "model":
|
||||
run_model(args.pred)
|
||||
run_model(args.pred, fuzzy=not args.strict)
|
||||
else:
|
||||
ap.print_help()
|
||||
|
||||
|
||||
16
watch_fetch.sh
Normal file
16
watch_fetch.sh
Normal file
@ -0,0 +1,16 @@
|
||||
#!/usr/bin/env bash
|
||||
# Live-Fortschritt des laufenden Prospekt-Downloads anzeigen.
|
||||
# bash watch_fetch.sh -> live mitlaufen (Strg-C zum Beenden)
|
||||
# bash watch_fetch.sh count -> nur kurzer Stand (X/435 + Groesse)
|
||||
LOG="/private/tmp/claude-501/-Users-florianherzog-development-AISE-Admin-Thesis--Florian-Code/745ff40e-0ecf-4430-886c-33fbc724df69/tasks/bp1wdlneq.output"
|
||||
DIR="data/rdf_poc/prose"
|
||||
|
||||
if [ "${1:-}" = "count" ]; then
|
||||
n=$(ls "$DIR"/*.txt 2>/dev/null | wc -l | tr -d ' ')
|
||||
last=$(grep -oE '\[[0-9]+/435\]' "$LOG" 2>/dev/null | tail -1)
|
||||
echo "Fertige Trusts: $n | letzte Logzeile: $last"
|
||||
du -sh "$DIR" 2>/dev/null
|
||||
else
|
||||
echo "Live-Download (Strg-C zum Beenden). Jede Zeile = ein Trust [n/435]:"
|
||||
tail -f "$LOG"
|
||||
fi
|
||||
Loading…
x
Reference in New Issue
Block a user