Upload 2 files

updated masks personnumber also better recognition of names.

anonymize.py

@@ -1,128 +1,141 @@
-import regex as re
-from typing import List, Tuple, Dict, Set
-from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
-
-
-class SwedishTextMasker:
-    def __init__(self, model_name: str = "RecordedFuture/Swedish-NER", threshold: float = 0.85):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForTokenClassification.from_pretrained(model_name)
-        self.ner_pipeline = pipeline("ner", model=self.model, tokenizer=self.tokenizer, aggregation_strategy="simple")
-        self.threshold = threshold
-
-    def _reconstruct_entities(self, tokens_with_labels: List[Tuple[str, str, float]]) -> List[Tuple[str, str]]:
-        words = []
-        current_word = ''
-        current_label = ''
-        scores = []
-
-        for token, label, score in tokens_with_labels:
-            if token.startswith('##'):
-                current_word += token[2:]
-                scores.append(score)
-            else:
-                if current_word:
-                    words.append((current_word, current_label , sum(scores) / len(scores)))
-                current_word, current_label  = token, label
-                scores = [score]
-
-        if current_word:
-            words.append((current_word, current_label , sum(scores) / len(scores)))
-
-        result = self._in_order_ent_list(words)
-        print("\n\n\n the result of inorder ent list : \n\n" , result)
-        return result
-
-
-    def _in_order_ent_list(self , all_ents_list):
-        threshold_ents = [ent for ent in all_ents_list if ent[2]>=self.threshold]
-        threshold_ents_word = {ent[0] for ent in threshold_ents}
-        result = [(ent[0] , ent[1]) for ent in all_ents_list if ent[0] in threshold_ents_word and len(ent[0]) >=2] 
-        return result
-    
-    def _get_chunks(self, text_list: List[str], chunk_size: int = 100) -> List[List[str]]:
-        return [text_list[i:i + chunk_size] for i in range(0, len(text_list), chunk_size)]
-
-    def _retrieve_ner(self, text: str) -> List[Tuple[str, str, float]]:
-        results = self.ner_pipeline(text)
-        return [
-            (ent["word"], ent["entity_group"], ent["score"])
-            for ent in results
-            if ent["entity_group"] in {"ORG", "PER", "TIT"}
-        ]
-
-    def _get_entities(self, text: str) -> List[Tuple[str, str, float]]:
-        tokens = text.split()
-        if len(tokens) > 100:
-            chunks = self._get_chunks(tokens)
-            all_ents = []
-            for chunk in chunks:
-                chunk_text = " ".join(chunk)
-                all_ents.extend(self._retrieve_ner(chunk_text))
-            return all_ents
-        else:
-            return self._retrieve_ner(text)
-
-    def _get_entity_dicts(self, entities: List[Tuple[str, str, float]]) -> Tuple[Dict[str, str], Dict[str, str], Set[str]]:
-        persons = [ent[0] for ent in entities if ent[1] == "PER"]
-        companies = [ent[0] for ent in entities if ent[1] == "ORG"]
-        titles = {ent[0] for ent in entities if ent[1] == "TIT"}
-
-        person_dict = {name: f"Person {chr(ord('A') + i)}" for i, name in enumerate(dict.fromkeys(persons))}
-        company_dict = {name: f"ORG_COMPANY {chr(ord('A') + i)}" for i, name in enumerate(dict.fromkeys(companies))}
-
-        return person_dict, company_dict, titles
-
-    @staticmethod
-    def mask_digits(text: str) -> str:
-        return re.sub(r'\d', 'x', text)
-
-    def mask_phone_numbers(self, text: str) -> str:
-        phone_regex = re.compile(r'(?:\+|00)?\d[\d\s\-()]{5,}\d')
-        return phone_regex.sub(lambda m: self.mask_digits(m.group()), text)
-
-    def mask_org_numbers(self, text: str) -> str:
-        org_regex = re.compile(r'\b\d{6}-?\d{4}\b')
-        return org_regex.sub(lambda m: self.mask_digits(m.group()), text)
-
-    def mask_emails(self, text: str) -> str:
-        email_regex = re.compile(r'\b([a-zA-Z0-9._%+-]+)@([a-zA-Z0-9.-]+\.[a-zA-Z]{2,})\b')
-        def email_masker(match):
-            local, domain = match.groups()
-            return f"{re.sub(r'[a-zA-Z0-9]', 'x', local)}@{re.sub(r'[a-zA-Z0-9]', 'x', domain)}"
-        return email_regex.sub(email_masker, text)
-
-    def mask_addresses(self, text: str) -> str:
-        address_regex = re.compile(
-            r'\b(?:[A-ZÅÄÖa-zåäöéÉèÈçÇß0-9\-]+\s)+\d{1,4}\s*,?\s*\d{3}\s?\d{2}\s+[A-ZÅÄÖa-zåäö\-]+',
-            re.UNICODE
-        )
-        return address_regex.sub('[ADDRESS]', text)
-
-    def mask_entities(self, text: str, entity_dict: Dict[str, str], tag: str) -> str:
-        for original, masked in entity_dict.items():
-            text = re.sub(re.escape(original), f"[{masked}]", text)
-        return text
-
-    def mask_titles(self, text: str, titles: Set[str]) -> str:
-        for title in titles:
-            text = re.sub(re.escape(title), "[Person_Title]", text)
-        return text
-
-    def mask_all(self, text: str) -> str:
-        old_text_backup = text
-        print("Original Text: \n\n" , text )
-        text = self.mask_phone_numbers(text)
-        text = self.mask_org_numbers(text)
-        text = self.mask_emails(text)
-        text = self.mask_addresses(text)
-
-        ents_raw = self._get_entities(old_text_backup)
-        ents = self._reconstruct_entities(ents_raw)
-        person_dict, company_dict, title_set = self._get_entity_dicts(ents)
-
-        text = self.mask_entities(text, company_dict, "ORG")
-        text = self.mask_entities(text, person_dict, "PER")
-        text = self.mask_titles(text, title_set)
-
+import regex as re
+from typing import List, Tuple, Dict, Set
+from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+
+
+class SwedishTextMasker:
+    def __init__(self, model_name: str = "RecordedFuture/Swedish-NER", threshold: float = 0.85):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForTokenClassification.from_pretrained(model_name)
+        self.ner_pipeline = pipeline("ner", model=self.model, tokenizer=self.tokenizer, aggregation_strategy="simple")
+        self.threshold = threshold
+
+    def _reconstruct_entities(self, tokens_with_labels: List[Tuple[str, str, float]]) -> List[Tuple[str, str]]:
+        words = []
+        current_word = ''
+        current_label = ''
+        scores = []
+        print("token with labels: " ,tokens_with_labels)
+        for token, label, score in tokens_with_labels:
+            if token.startswith('##'):
+                current_word += token[2:]
+                scores.append(score)
+            else:
+                if current_word:
+                    words.append((current_word, current_label , sum(scores) / len(scores)))
+                current_word, current_label  = token, label
+                scores = [score]
+
+        if current_word:
+            words.append((current_word, current_label , sum(scores) / len(scores)))
+
+        words = [(re.sub(r'\s*-\s*', '-', w[0]), w[1], w[2]) for w in words]
+        result = self._in_order_ent_list(words)
+        print("\n\n\n the result of inorder ent list : \n\n" , result)
+        return result
+
+
+    def _in_order_ent_list(self , all_ents_list):
+        threshold_ents = [ent for ent in all_ents_list if ent[2]>=self.threshold]
+        threshold_ents_word = {ent[0] for ent in threshold_ents}
+        result = [(ent[0] , ent[1]) for ent in all_ents_list if ent[0] in threshold_ents_word and len(ent[0]) >=2] 
+        return result
+    
+    def _get_chunks(self, text_list: List[str], chunk_size: int = 100) -> List[List[str]]:
+        return [text_list[i:i + chunk_size] for i in range(0, len(text_list), chunk_size)]
+
+    def _retrieve_ner(self, text: str) -> List[Tuple[str, str, float]]:
+        results = self.ner_pipeline(text)
+        return [
+            (ent["word"], ent["entity_group"], ent["score"])
+            for ent in results
+            if ent["entity_group"] in {"ORG", "PER", "TIT"}
+        ]
+
+    def _get_entities(self, text: str) -> List[Tuple[str, str, float]]:
+        tokens = text.split()
+        if len(tokens) > 100:
+            chunks = self._get_chunks(tokens)
+            all_ents = []
+            for chunk in chunks:
+                chunk_text = " ".join(chunk)
+                all_ents.extend(self._retrieve_ner(chunk_text))
+            return all_ents
+        else:
+            return self._retrieve_ner(text)
+
+    def _get_entity_dicts(self, entities: List[Tuple[str, str, float]]) -> Tuple[Dict[str, str], Dict[str, str], Set[str]]:
+        persons = [ent[0] for ent in entities if ent[1] == "PER"]
+        companies = [ent[0] for ent in entities if ent[1] == "ORG"]
+        titles = {ent[0] for ent in entities if ent[1] == "TIT"}
+
+        person_dict = {name: f"Person {chr(ord('A') + i)}" for i, name in enumerate(dict.fromkeys(persons))}
+        company_dict = {name: f"ORG_COMPANY {chr(ord('A') + i)}" for i, name in enumerate(dict.fromkeys(companies))}
+
+        return person_dict, company_dict, titles
+
+
+    def remove_personnummer(self, text):
+        pattern = r"""
+        (?<!\d)
+        (?:
+            \d{6,8}[\-\s]?\d{4} |
+            \d{5}[\-\s]?\d{4}
+        )
+        (?!\d)
+        """
+        return re.sub(pattern, '[person/org nummer]', text, flags=re.VERBOSE)
+    
+    @staticmethod
+    def mask_digits(text: str) -> str:
+        return re.sub(r'\d', 'x', text)
+
+    def mask_phone_numbers(self, text: str) -> str:
+        phone_regex = re.compile(r'(?:\+|00)?\d[\d\s\-()]{5,}\d')
+        return phone_regex.sub(lambda m: self.mask_digits(m.group()), text)
+
+    # def mask_org_numbers(self, text: str) -> str:
+    #     org_regex = re.compile(r'\b\d{6}-?\d{4}\b')
+    #     return org_regex.sub(lambda m: self.mask_digits(m.group()), text)
+
+    def mask_emails(self, text: str) -> str:
+        email_regex = re.compile(r'\b([a-zA-Z0-9._%+-]+)@([a-zA-Z0-9.-]+\.[a-zA-Z]{2,})\b')
+        def email_masker(match):
+            local, domain = match.groups()
+            return f"{re.sub(r'[a-zA-Z0-9]', 'x', local)}@{re.sub(r'[a-zA-Z0-9]', 'x', domain)}"
+        return email_regex.sub(email_masker, text)
+
+    def mask_addresses(self, text: str) -> str:
+        address_regex = re.compile(
+            r'\b(?:[A-ZÅÄÖa-zåäöéÉèÈçÇß0-9\-]+\s)+\d{1,4}\s*,?\s*\d{3}\s?\d{2}\s+[A-ZÅÄÖa-zåäö\-]+',
+            re.UNICODE
+        )
+        return address_regex.sub('[ADDRESS]', text)
+
+    def mask_entities(self, text: str, entity_dict: Dict[str, str], tag: str) -> str:
+        for original, masked in entity_dict.items():
+            text = re.sub(re.escape(original), f"[{masked}]", text)
+        return text
+
+    def mask_titles(self, text: str, titles: Set[str]) -> str:
+        for title in titles:
+            text = re.sub(re.escape(title), "[Person_Title]", text)
+        return text
+
+    def mask_all(self, text: str) -> str:
+        old_text_backup = text
+        text = self.remove_personnummer(text)
+        text = self.mask_phone_numbers(text)
+        # text = self.mask_org_numbers(text)
+        text = self.mask_emails(text)
+        text = self.mask_addresses(text)
+
+        ents_raw = self._get_entities(old_text_backup)
+        ents = self._reconstruct_entities(ents_raw)
+        person_dict, company_dict, title_set = self._get_entity_dicts(ents)
+
+        text = self.mask_entities(text, company_dict, "ORG")
+        text = self.mask_entities(text, person_dict, "PER")
+        text = self.mask_titles(text, title_set)
+
         return text

app.py

@@ -1,33 +1,64 @@
-import fitz  # PyMuPDF
-import gradio as gr
-from anonymize import SwedishTextMasker
-
-# Instantiate once, globally
-text_anonymizer = SwedishTextMasker(threshold= 0.9)
-
-def extract_text_from_pdf(pdf_file):
-    if pdf_file is None:
-        return "No file uploa9999ded."
-
-    # Approach 1: open via file path (usually safer)
-    with fitz.open(pdf_file.name) as doc:
-        text_output = ""
-        for page in doc:
-            text_output += page.get_text()
-
-    raw_text = text_output.strip()
-    anonymized_text = text_anonymizer.mask_all(raw_text)
-    return anonymized_text
-
-# Gradio interface
-with gr.Blocks(title="PDF -> Anonymized Text") as demo:
-    gr.Markdown("### 📄 PDF Anonymizer (text only, skips images)")
-    with gr.Row():
-        pdf_input = gr.File(label="Upload a PDF", file_types=[".pdf"])
-    text_output = gr.Textbox(label="Anonymized Output", lines=20, interactive=False)
-
-    extract_button = gr.Button("Anonymize Text")
-    extract_button.click(fn=extract_text_from_pdf, inputs=pdf_input, outputs=text_output)
-
-if __name__ == "__main__":
-    demo.launch()
+import fitz  # PyMuPDF
+import gradio as gr
+from anonymize import SwedishTextMasker
+
+# Instantiate once, globally
+text_anonymizer = SwedishTextMasker(threshold= 0.5)
+
+def join_short_lines(text, min_length=30):
+    """
+    Joins lines that are shorter than min_length with the next line.
+    """
+    lines = text.split('\n')
+    new_lines = []
+    buffer = ""
+    for line in lines:
+        stripped = line.strip()
+        if not stripped:
+            if buffer:
+                new_lines.append(buffer)
+                buffer = ""
+            new_lines.append("")  # preserve empty lines
+            continue
+        if len(stripped) < min_length and not stripped.endswith(('.', ':', ';', '?', '!')):
+            buffer += " " + stripped if buffer else stripped
+        else:
+            if buffer:
+                buffer += " " + stripped
+                new_lines.append(buffer)
+                buffer = ""
+            else:
+                new_lines.append(stripped)
+    if buffer:
+        new_lines.append(buffer)
+    return "\n".join(new_lines)
+
+
+def extract_text_from_pdf(pdf_file):
+    if pdf_file is None:
+        return "No file uploaded."
+
+    # Approach 1: open via file path (usually safer)
+    with fitz.open(pdf_file.name) as doc:
+        text_output = ""
+        for page in doc:
+            text_output += page.get_text(flags=1)
+
+    raw_text = text_output.strip()
+    # raw_text = join_short_lines(raw_text)  # <--- Add this line!
+    print(raw_text)
+    anonymized_text = text_anonymizer.mask_all(raw_text)
+    return anonymized_text
+
+# Gradio interface
+with gr.Blocks(title="PDF -> Anonymized Text") as demo:
+    gr.Markdown("### 📄 PDF Anonymizer (text only, skips images)")
+    with gr.Row():
+        pdf_input = gr.File(label="Upload a PDF", file_types=[".pdf"])
+    text_output = gr.Textbox(label="Anonymized Output", lines=20, interactive=False)
+
+    extract_button = gr.Button("Anonymize Text")
+    extract_button.click(fn=extract_text_from_pdf, inputs=pdf_input, outputs=text_output)
+
+if __name__ == "__main__":
+    demo.launch()