hiteshwadhwani/pii-model-indicv2

from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

model = AutoModelForTokenClassification.from_pretrained("hiteshwadhwani/pii-model-indicv2")
tokenizer = AutoTokenizer.from_pretrained("hiteshwadhwani/pii-model-indicv2")

ner = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="first")

# English
results = ner("Mr John Smith lives at 456 Oak Avenue Boston")

# Hinglish
results = ner("Mera naam Rajesh Kumar hai aur main Mumbai mein rehta hoon")

# Hindi (Devanagari)
results = ner("मेरा नाम राजेश कुमार है और मैं मुंबई में रहता हूं")

# Hindi amounts and dates in words
results = ner("आपके क्रेडिट कार्ड के बीस हज़ार रूपये शुल्क लंबित हैं, जो पंद्रह मार्च को देना था।")

for entity in results:
    if entity["entity_group"] != "O":
        print(f"{entity['word']} → {entity['entity_group']} ({entity['score']:.2f})")

def redact_pii(text, ner_pipeline, threshold=0.5):
    results = ner_pipeline(text)
    entities = [r for r in results if r["score"] >= threshold and r["entity_group"] != "O"]
    entities.sort(key=lambda x: x["start"])

    merged = []
    for ent in entities:
        label = ent["entity_group"]
        if merged and merged[-1]["label"] == label and ent["start"] <= merged[-1]["end"] + 1:
            merged[-1]["end"] = max(merged[-1]["end"], ent["end"])
        else:
            merged.append({"label": label, "start": ent["start"], "end": ent["end"]})

    redacted = text
    for span in reversed(merged):
        redacted = redacted[:span["start"]] + f"[{span['label']}]" + redacted[span["end"]:]
    return redacted

print(redact_pii("Shri Rajesh Kumar lives at 42 MG Road Bengaluru Karnataka", ner))
# [PREFIX] [FIRSTNAME] [LASTNAME] lives at [BUILDINGNUMBER] [STREET] [CITY] [STATE]

print(redact_pii("आपके बीस हज़ार रूपये शुल्क लंबित हैं, जो पंद्रह मार्च को देना था।", ner))
# आपके [AMOUNT] शुल्क लंबित हैं, जो [DATE] को देना था।