Upload bias_detector.py

bias_detector.py

@@ -18,7 +18,6 @@
 
 import re
 import torch
-import spacy
 from collections import Counter
 from transformers import (
     pipeline,
@@ -59,9 +58,8 @@ class BiasDetector:
         Initialize all models. Call once at startup, reuse for every request.
 
         Args:
-            llm_model_name: HuggingFace instruction-tuned model for rewriting
-                            and CV anonymization. Defaults to Qwen2.5-7B-Instruct
-                            loaded in 4-bit to fit on CPU-only Spaces (~4.5 GB RAM).
+            llm_model_name: HuggingFace instruction-tuned model for rewriting.
+                            Defaults to Qwen2.5-1.5B-Instruct (float16, CPU).
         """
         print("Loading bias detection model...")
         self.classifier = pipeline(
@@ -70,8 +68,13 @@ class BiasDetector:
             device=0 if torch.cuda.is_available() else -1,
         )
 
-        print("Loading spaCy NER model...")
-        self.nlp = spacy.load("en_core_web_sm")
+        print("Loading transformer NER model...")
+        self.ner = pipeline(
+            "ner",
+            model="dslim/bert-base-NER",
+            aggregation_strategy="simple",
+            device=0 if torch.cuda.is_available() else -1,
+        )
 
         print(f"Loading {llm_model_name}...")
         self.tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
@@ -418,7 +421,7 @@ Output: "They pursue excellence and they work hard. They have delivered results.
     def _surface_anonymize(self, text: str) -> str:
         """
         Step A: Replace emails with [EMAIL].
-        Step B: spaCy NER → frequency-based person labelling.
+        Step B: Transformer NER (dslim/bert-base-NER) → frequency-based person labelling.
                 Most-mentioned person → [CANDIDATE].
                 Others → [PERSON_2], [PERSON_3], … in order of first appearance.
         Step C: Rule-based pronoun / title / gendered-noun replacement.
@@ -442,24 +445,22 @@ Output: "They pursue excellence and they work hard. They have delivered results.
             text,
         )
 
-        # --- Step B: Person names via spaCy NER ---
-        doc = self.nlp(anonymized)
-
-        # Collect all PERSON spans and count mention frequency per canonical name
-        # (use the first token as a rough canonical key to handle "Sarah" vs "Sarah Johnson")
-        # Academic degree keywords — spaCy sometimes mislabels degree names as PERSON
-        # e.g. "MSc Machine Learning", "BSc Mathematics", "PhD Computer Science"
-        DEGREE_KEYWORDS = {
-            "msc", "bsc", "ba", "ma", "mba", "phd", "llb", "llm", "beng", "meng",
-            "doctorate", "bachelor", "master", "masters", "graduate", "postgraduate",
-        }
-
-        person_spans = [
-            (ent.start_char, ent.end_char, ent.text)
-            for ent in doc.ents
-            if ent.label_ == "PERSON"
-            and not any(token.lower() in DEGREE_KEYWORDS for token in ent.text.split())
-        ]
+        # --- Step B: Person names via transformer NER (dslim/bert-base-NER) ---
+        # Returns entity_group PER/ORG/LOC/MISC — we only keep PER.
+        # Unlike spaCy en_core_web_sm, this model correctly distinguishes
+        # company names (ORG) and locations (LOC) from person names (PER),
+        # eliminating false positives like "Luminary Analytics" or "Machine Learning".
+        ner_results = self.ner(anonymized)
+        person_spans = []
+        for ent in ner_results:
+            if ent["entity_group"] != "PER":
+                continue
+            # Clean BERT subword artifacts (## prefixes from tokenizer)
+            word = ent["word"].replace("##", "").strip()
+            # Skip if too short to be a real name token (avoids partial matches)
+            if len(word) < 3:
+                continue
+            person_spans.append((ent["start"], ent["end"], word))
 
         # Count frequency by normalised name (lower-case first token)
         name_freq: Counter = Counter()
@@ -470,7 +471,7 @@ Output: "They pursue excellence and they work hard. They have delivered results.
             if key not in first_seen:
                 first_seen[key] = start
 
-        # Seed name_freq with email-derived names not already found by spaCy
+        # Seed name_freq with email-derived names not already found by NER
         for first, last, full in email_names:
             key = first.lower()
             if key not in name_freq:
@@ -491,36 +492,60 @@ Output: "They pursue excellence and they work hard. They have delivered results.
             for i, k in enumerate(other_keys, start=2):
                 label_map[k] = f"[PERSON_{i}]"
 
-            # Replace NER spans in reverse order to preserve char offsets
-            for start, end, name in reversed(person_spans):
+            # Extend person spans: if the token immediately after a PER span
+            # is a capitalised word not in common vocab, treat it as a surname
+            # e.g. NER finds "Marcus" but misses "Obi" → extend to "Marcus Obi"
+            extended_spans = []
+            for start, end, name in person_spans:
                 key = name.strip().lower().split()[0]
                 label = label_map.get(key, "[PERSON]")
+                # Check if next token after span is a capitalised word (surname candidate)
+                rest = anonymized[end:]
+                surname_match = re.match(r"^\s+([A-Z][A-Za-z]{1,20})\b", rest)
+                if surname_match:
+                    candidate_surname = surname_match.group(1)
+                    # Only extend if it's not a common non-name word
+                    NON_NAMES = {"The", "This", "That", "Their", "They", "He", "She",
+                                 "His", "Her", "During", "In", "At", "For", "And",
+                                 "But", "With", "From", "To", "Of", "On", "By"}
+                    if candidate_surname not in NON_NAMES:
+                        new_end = end + len(surname_match.group(0))
+                        new_name = anonymized[start:new_end].strip()
+                        extended_spans.append((start, new_end, new_name, label))
+                        # Also register the surname token in full_name_label later
+                        continue
+                extended_spans.append((start, end, name, label))
+
+            # Replace NER spans in reverse order to preserve char offsets
+            for start, end, name, label in reversed(extended_spans):
                 anonymized = anonymized[:start] + label + anonymized[end:]
 
             # Second pass: regex sweep for any remaining occurrences of known names
             # (catches names spaCy missed because they were next to org/title context).
             # Sort longest-first to avoid partial replacements.
             full_name_label: dict = {}
-            for start, end, name in person_spans:
-                key = name.strip().lower().split()[0]
-                label = label_map.get(key, "[PERSON]")
+            for start, end, name, label in extended_spans:
                 full_name_label[name.strip()] = label
-            # Also add email-derived names that spaCy missed
+                # Also register individual tokens (first name, last name separately)
+                for token in name.strip().split():
+                    if len(token) >= 4:
+                        full_name_label.setdefault(token, label)
+            # Also add email-derived names that NER missed
             for first, last, full in email_names:
                 key = first.lower()
                 if key not in label_map:
                     continue
-                # Add both full name and individual tokens so all forms are caught
+                # Add full name, first name, and last name so all forms are caught
                 for variant in [full, first, last]:
                     if variant not in full_name_label:
                         full_name_label[variant] = label_map[key]
 
-            # Also add individual tokens (first name, last name) for each known person
-            # so partial matches like "[CANDIDATE] Obi" get cleaned up
+            # Build token_label AFTER full_name_label is complete (including email names)
+            # so last names like "Obi" or "Reeves" from email are included
             token_label: dict = {}
             for full_name, label in full_name_label.items():
                 for token in full_name.split():
-                    if len(token) > 2 and token not in token_label:
+                    if len(token) >= 4 and token not in token_label:
                         token_label[token] = label
 
             # Replace full names first (longest first), then individual tokens