Team Ai
Apppublic

gunasekhar2003/Email_Classification

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
utils.py73 linesDownload Raw Back to root
1import re2import spacy3import subprocess4 5def ensure_spacy_model():6    try:7        spacy.load("en_core_web_sm")8    except OSError:9        print("Downloading spaCy model...")10        subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])11 12ensure_spacy_model()13nlp = spacy.load("en_core_web_sm")14 15def detect_regex_entities(text):16    entities = []17 18    patterns = {19        'email': r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',20        'phone': r'\b\d{10,12}\b',21        'aadhar': r'\b\d{4}\s\d{4}\s\d{4}\b',22        'expiry': r'\b(0[1-9]|1[0-2])\/?([0-9]{2})\b',23        'cvv': r'\b[0-9]{3}\b',24    }25 26    for label, pattern in patterns.items():27        for match in re.finditer(pattern, text):28            entities.append({29                'start': match.start(),30                'end': match.end(),31                'entity': match.group(),32                'label': label33            })34 35    return entities36 37def detect_name_entities(text):38    entities = []39    doc = nlp(text)40    for ent in doc.ents:41        if ent.label_ == "PERSON":42            entities.append({43                'start': ent.start_char,44                'end': ent.end_char,45                'entity': ent.text,46                'label': 'name'47            })48    return entities49 50def mask_pii(text):51    regex_entities = detect_regex_entities(text)52    name_entities = detect_name_entities(text)53    all_entities = regex_entities + name_entities54 55   56    entities_sorted = sorted(all_entities, key=lambda x: x['start'])57 58    masked_text = ""59    current_pos = 060    for ent in entities_sorted:61        start, end = ent['start'], ent['end']62        label = ent['label'].upper()63 64        if start >= current_pos:  65            masked_text += text[current_pos:start] + f"<{label} MASKED>"66            current_pos = end67 68    69    masked_text += text[current_pos:]70 71    return masked_text, all_entities72 73