aravagarwal/CodeCloakPII

1from tree_sitter import Language2from tree_sitter import Parser3import json4import os5import re6import warnings7from typing import Optional as Opt8from collections import Counter9 10 11from presidio_analyzer import (12    AnalyzerEngine,13    BatchAnalyzerEngine,14    Pattern,15    PatternRecognizer,16)17from presidio_analyzer.nlp_engine import TransformersNlpEngine, NlpEngineProvider18from presidio_anonymizer import AnonymizerEngine, BatchAnonymizerEngine19from presidio_anonymizer.entities import OperatorResult, OperatorConfig20from presidio_analyzer import PatternRecognizer21 22from tqdm import tqdm23import typer24import spacy25 26# Putting package-level imports here as a test - probably terrible27 28Language.build_library(29    # Store the library in the `build` directory30    "build/my-languages.so",31    # Include one or more languages32    ["lib/tree-sitter-java", "lib/tree-sitter-python"],33)34 35JAVA_LANGUAGE = Language("build/my-languages.so", "java")36PYTHON_LANGUAGE = Language("build/my-languages.so", "python")37 38 39parser = Parser()40parser.set_language(PYTHON_LANGUAGE)41 42configuration = {43    "nlp_engine_name": "spacy",44    "models": [{"lang_code": "en", "model_name": "en_core_web_trf"}],45    "ner_model_configuration": {46        "labels_to_ignore": ["O"],47        "aggregation_strategy": "simple",48        "stride": 16,49        "alignment_mode": "expand",50        "model_to_presidio_entity_mapping": {51            "PER": "PERSON",52            "PERSON": "PERSON",53            "LOC": "LOCATION",54            "ORG": "ORGANIZATION",55            "AGE": "AGE",56            "ID": "ID",57            "EMAIL": "EMAIL",58            "PATIENT": "PERSON",59            "STAFF": "PERSON",60            "HOSP": "ORGANIZATION",61            "PATORG": "ORGANIZATION",62            "DATE": "DATE_TIME",63            "PHONE": "PHONE_NUMBER",64            "HCW": "PERSON",65            "HOSPITAL": "ORGANIZATION",66            "FIRSTNAME": "PERSON",67            "FULLNAME": "PERSON",68            "NAME": "PERSON",69            "GENDER": "PERSON",70            "DISPLAYNAME": "PERSON",71            "MIDDLENAME": "PERSON",72            "ACCOUNTNAME": "PERSON",73            "CARDINAL": "AGE",74            "FAC": "AGE",75            "PRODUCT": "AGE",76            "WORK_OF_ART": "AGE",77            "QUANTITY": "AGE",78            "LANGUAGE": "AGE",79            "PERCENT": "AGE",80            "TIME": "AGE",81            "ORDINAL": "AGE",82            "MONEY": "AGE",83            "GPE": "LOCATION",84            "NORP": "PERSON",85        },86    },87}88 89MAX_LENGTH = 100000090MAX_STRINGS_IN_SAMPLE = 10091 92provider = NlpEngineProvider(nlp_configuration=configuration)93provider = provider.create_engine()94provider.nlp["en"].max_length = MAX_LENGTH95 96analyzer = AnalyzerEngine(nlp_engine=provider)97anonymizer = AnonymizerEngine()98 99regex = r'(?!https?://)(?:[a-zA-Z]:\\(?:[^\\/:*?"<>|\r\n]+\\)+[^\\/:*?"<>|\s\r\n]*|\/(?:[^\\/:*?"<>|\r\n]+\/[^\\/:*?"<>|\s\r\n]+\/)*[^\\/:*?"<>|\s\r\n]*)'100path_pattern = Pattern(name="path_pattern", regex=regex, score=1.0)101 102path_re = PatternRecognizer(supported_entity="URL", patterns=[path_pattern])103analyzer.registry.add_recognizer(path_re)104 105# Not needed - just metrics106 107REMOVED_SET = Counter()108 109