feat: adding regex and transformer

.gitignore

@@ -0,0 +1 @@
+Pipfile

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
+from lib.masker.masker import run_mask
 
-def greet(name):
-    return "Hello " + name + "!!"
+def mask(text, options):
+    return run_mask(text, options)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+demo = gr.Interface(fn=mask, inputs=["text", "text"], outputs="text")
+demo.launch(debug=True)

lib/masker/masker.py

@@ -0,0 +1,18 @@
+from lib.masker.regex import REGEX_FUNCTIONS
+from lib.masker.transformer import TRANSFORMER_FUNCTIONS
+
+DEFAULT_OPTIONS = [
+    "name", "email", "phone", "credit_card", "local"
+]
+
+def run_mask(text, options):
+    if not isinstance(options, list) or not options:
+        options = DEFAULT_OPTIONS
+
+    for option in options:
+        if option in TRANSFORMER_FUNCTIONS:
+            text = TRANSFORMER_FUNCTIONS[option](text)
+        if option in REGEX_FUNCTIONS:
+            text = REGEX_FUNCTIONS[option](text)
+
+    return text

lib/masker/ner.py

@@ -0,0 +1,25 @@
+from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
+
+entity_types = {
+    'B-PESSOA': 'name',
+    'I-PESSOA': 'name',
+    'B-LOCAL': 'local',
+    'I-LOCAL': 'local',
+}
+
+def ner_exec(text):
+    model_name = "liaad/NER_harem_bert-base-portuguese-cased"
+    model = AutoModelForTokenClassification.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
+    return ner_pipeline(text)
+
+def ner_filter(ner_results, text, mask_type):
+    for entity in ner_results:
+        entity_word = entity['word']
+        type_name = entity_types[entity['entity']]
+
+        if (type_name == mask_type):
+            text = text.replace(entity_word, "*" * len(entity_word))
+
+    return text

lib/masker/regex.py

@@ -0,0 +1,16 @@
+import re
+
+def mask_credit_card(text):
+    return re.sub(r'(\d{12})\d{4}', r'\1****', text)
+
+def mask_phone(text):
+    return re.sub(r'\+55\s?(\d{2})\s?(\d{5})[-\s]?(\d{4})', r'(\1) XXX-XXXX', text)
+
+def mask_email(text):
+    return re.sub(r'([a-zA-Z0-9._%+-])([a-zA-Z0-9.-]+)@([a-zA-Z0-9.-]+)', r'\1*****@\3', text)
+
+REGEX_FUNCTIONS = {
+    'credit_card': mask_credit_card,
+    'phone': mask_phone,
+    'email': mask_email,
+}

lib/masker/transformer.py

@@ -0,0 +1,18 @@
+from lib.masker.ner import ner_exec, ner_filter
+    
+def mask_name(text):
+    mask_type = "name"
+    ner_results = ner_exec(text)
+
+    return ner_filter(ner_results, text, mask_type)
+
+def mask_local(text):
+    mask_type = "local"
+    ner_results = ner_exec(text)
+
+    return ner_filter(ner_results, text, mask_type)
+
+TRANSFORMER_FUNCTIONS = {
+    'name': mask_name,
+    'local': mask_local,
+}

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+transformers
+torch