rodado
891d58ec55
Substitui o gitGraph do Mermaid (ilegível com 3.437 commits) por uma timeline interativa com eixo temporal real, filtros por instituição, tipo de documento e pessoas mencionadas (extraidas via NER local com pierreguillou/ner-bert-base-cased-pt-lenerbr).
55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
import json, re, time
|
|
from collections import Counter, defaultdict
|
|
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
|
|
|
|
docs = json.load(open("docs.json"))
|
|
texts = [d["text"] for d in docs]
|
|
print(f"{len(texts)} documents to process", flush=True)
|
|
|
|
name = "pierreguillou/ner-bert-base-cased-pt-lenerbr"
|
|
tok = AutoTokenizer.from_pretrained(name)
|
|
model = AutoModelForTokenClassification.from_pretrained(name)
|
|
nlp = pipeline("ner", model=model, tokenizer=tok, aggregation_strategy="simple", device=-1)
|
|
|
|
def merge_entities(ents, group="PESSOA"):
|
|
ents = [e for e in ents if e["entity_group"] == group]
|
|
ents.sort(key=lambda e: e["start"])
|
|
merged = []
|
|
for e in ents:
|
|
if merged and e["start"] - merged[-1]["end"] <= 1:
|
|
merged[-1]["end"] = e["end"]
|
|
else:
|
|
merged.append({"start": e["start"], "end": e["end"]})
|
|
return merged
|
|
|
|
t0 = time.time()
|
|
BATCH = 32
|
|
all_people_per_doc = [] # list aligned with docs: list of names
|
|
for i in range(0, len(texts), BATCH):
|
|
batch = texts[i:i+BATCH]
|
|
results = nlp(batch)
|
|
for text, ents in zip(batch, results):
|
|
merged = merge_entities(ents, "PESSOA")
|
|
names = [text[m["start"]:m["end"]].strip() for m in merged]
|
|
names = [n for n in names if len(n) >= 3]
|
|
all_people_per_doc.append(names)
|
|
if (i//BATCH) % 10 == 0:
|
|
elapsed = time.time()-t0
|
|
print(f"{i+len(batch)}/{len(texts)} docs, {elapsed:.1f}s elapsed", flush=True)
|
|
|
|
print("total time", time.time()-t0, "s")
|
|
|
|
# save raw per-doc results
|
|
out = []
|
|
for d, names in zip(docs, all_people_per_doc):
|
|
out.append({"id": d["id"], "names": names})
|
|
json.dump(out, open("ner_people_per_doc.json","w",encoding="utf-8"), ensure_ascii=False)
|
|
|
|
counter = Counter()
|
|
for names in all_people_per_doc:
|
|
for n in names:
|
|
counter[n] += 1
|
|
print("distinct raw person strings:", len(counter))
|
|
for n,c in counter.most_common(50):
|
|
print(c, n)
|