src/redacto/client/engine.py

import torch
import pandas as pd
from transformers import pipeline
from tqdm import tqdm  # For a progress bar
import re
from gliner import GLiNER

class UniversalRedactionEngine:
    def __init__(self):
        print("⚡ Loading AI Models... (This uses GPU if available)")

        device = "cuda" if torch.cuda.is_available() else "cpu"
        # 1. NER Model (The "Finder")
        # Finds names/orgs/locations with high recall.
        self.gliner_model = GLiNER.from_pretrained("nvidia/gliner-pii").to(device)
        self.target_labels = [
            "email", "phone_number", "ssn", "date_of_birth", "date", "street_address", "account_number",
            "user_name", "company_name",
        ]
        
        # 2. Zero-Shot Classifier (The "Judge")
        # Decides context without hard-coded rules.
        self.classifier = pipeline(
            "zero-shot-classification",
            model="facebook/bart-large-mnli",
            device=0 if torch.cuda.is_available() else -1
        )
        
        # Define the semantic concepts (You can change these to fit other domains)
        self.sensitivity_labels = ["public official or government entity", "private individual or entity"]



    def _assess_sensitivity(self, entity_text, context_sentence):
        """
        Universal Logic: Checks if the entity is sensitive regardless of its type.
        """
        # Hypothesis: "The entity [X] represents [LABEL]."
        result = self.classifier(
            context_sentence,
            candidate_labels=self.sensitivity_labels,
            hypothesis_template=f"The entity '{entity_text}' represents {{}}."
        )
        
        # Return the winning label and its score
        return result['labels'][0], result['scores'][0]
        

    def process_dataframe(self, df_layout):
        """
        Consumes the Phase 1 DataFrame and returns a Redaction Plan.
        """
        redaction_candidates = []
        
        print(f"🕵️ Scanning {len(df_layout)} layout blocks for PII...")

        # Iterate through the DataFrame with a progress bar
        for index, row in tqdm(df_layout.iterrows(), total=df_layout.shape[0]):
            
            text_block = str(row.get('text', ''))
            page_num = row.get('page', 1)
            bbox = row.get('bbox', None)
            label = row.get('label', 'text')
            
            # Extract the new context headers!
            h1 = row.get('context_h1', '')
            h2 = row.get('context_h2', '')
            
            if len(text_block) < 5: continue

            # --- THE MAGIC UPGRADE: Build the Enriched Context ---
            # We tell BART exactly where this text lives in the document structure.
            if label == 'table':
                enriched_context = f"Under the section '{h1} > {h2}', the following data table is presented:\n{text_block}"
            else:
                enriched_context = f"Under the section '{h1} > {h2}', the document states: {text_block}"


           # --- PHASE 2: AI ENTITY MATCHING (Context PII) ---
            entities = self.gliner_model.predict_entities(text_block, self.target_labels, threshold=0.5)
            
            for ent in entities:
                entity_text = ent['text']
                entity_type = ent['label'] 

                if entity_type in ["email", "phone_number", "ssn", "date_of_birth", "street_address", "account_number"]:
                    redaction_candidates.append({
                        "page": page_num,
                        "text_found": entity_text,
                        "entity_type": entity_type,
                        "decision": "REDACT",
                        "reason": "Strict PII (GLiNER)",
                        "confidence": 0.99,
                        "bbox": bbox,
                        "full_text": text_block
                    })
                    continue
                
                # --- STEP 2: DECIDE SENSITIVITY ---
                # We ask the AI: "Is this Public or Private?"
                verdict, confidence = self._assess_sensitivity(entity_text, text_block)
                
                # --- STEP 3: ACT ---
                # If AI says "Sensitive Private Information", we redact.
                if verdict == "private individual or entity" and confidence > 0.60:
                    redaction_candidates.append({
                        "page": page_num,
                        "text_found": entity_text,
                        "entity_type": entity_type,
                        "decision": "REDACT",
                        "reason": verdict,
                        "confidence": confidence,
                        "bbox": bbox,
                        "full_text": text_block
                    })

                else:
                    # 🛡️ THE AUDIT LOG: Prove the AI is filtering out public entities
                    print(f"  -> 🛡️ Context Judge Ignored: '{entity_text}' | Reason: {verdict} ({confidence:.1%} conf)")
            # Else: It is "Public Information" -> Do nothing.

        return pd.DataFrame(redaction_candidates)