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)