metricspace/GDPR_Input_Detection_and_Anonymization_0.5B

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_path = "metricspace/GDPR_Input_Detection_and_Anonymization_0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_path)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)

import re
import json

# Set tokenizer tokens and model padding
tokenizer.pad_token = "<|im_start|>"
tokenizer.eos_token = "<|im_end|>"
tokenizer.padding_side = "left"
model.generation_config.pad_token_id = tokenizer.pad_token_id

# Define formats for different analysis types
formats = {
    "sensitivity": """<|im_start|>system\nSensitivity<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n""",
    "complexity": """<|im_start|>system\nComplexity<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n""",
    "entity_detection": """<|im_start|>system\nEntity Detection<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n""",
    "entity_swapping": """<|im_start|>system\nEntity Swapping<|im_end|>\n<|im_start|>user\nentities:\n{entities}\ntext:\n{text}<|im_end|>\n<|im_start|>assistant\n"""
}

def model_inference(text, mode="anonymization", max_new_tokens=2028, config=None, entity_mapping=None, return_entities=False, reverse_mapping=False):
    if mode not in formats and mode != "anonymization":
        raise ValueError("Invalid mode. Choose from 'sensitivity', 'complexity', 'entity_detection', 'anonymization'.")

    
    # Configuration for anonymization
    # The `config` dictionary specifies the anonymization behavior for each type of entity detected.
    # Each key in `config` represents an entity type (e.g., "LOC" for location, "PERSON" for personal names),
    # and the value assigned to that key determines how entities of that type should be anonymized:
    # 
    # - "RANDOM": Replaces the entity with a randomly selected placeholder.
    # - "GENERAL LOW", "GENERAL MEDIUM", "GENERAL HIGH": Replaces the entity with a generalized label,
    #   with the intensity level (LOW, MEDIUM, HIGH) controlling the specificity. For example, 
    #   "GENERAL LOW" might use a more specific label ("Local Park") while "GENERAL HIGH" would use
    #   a broader label ("Recreational Area").
    # 
    # This allows fine-grained control over anonymization, ensuring that different types of sensitive 
    # information can be replaced in ways that are appropriate for the context. For example:
    #   - "LOC": "RANDOM" replaces any detected location with a random placeholder.
    #   - "DATETIME": "GENERAL LOW" uses a lower-intensity generalization for dates and times.
    # 
    # This flexibility enables custom anonymization policies to suit different privacy or obfuscation needs.

    if config is None:
        config = {
            "LOC": "RANDOM",
            "PERSON": "RANDOM",
            "DEM": "RANDOM",
            "CODE": "RANDOM",
            "ORG": "GENERAL MEDIUM",
            "DATETIME": "GENERAL LOW",
            "QUANTITY": "RANDOM",
            "MISC": "RANDOM",
        }

    # Anonymization Mode
    if mode == "anonymization":
        # Step 1: Entity detection
        detection_prompt = formats["entity_detection"].format(text=text)
        detection_inputs = tokenizer(detection_prompt, return_tensors="pt").to(device)
        detection_output = model.generate(
            **detection_inputs,
            max_new_tokens=max_new_tokens,
            use_cache=True,
            eos_token_id=151645
        )
        detection_text = tokenizer.decode(detection_output[0], skip_special_tokens=True)
        detected_entities = postprocess_entity_recognition(detection_text)
        
        # Step 2: Select entities based on config
        selected_entities = select_entities_based_on_json(detected_entities, config)
        entities_str = "\n".join([f"{entity} : {label}" for entity, label in selected_entities])
        # Step 3: Entity swapping for anonymization
        swapping_prompt = formats["entity_swapping"].format(entities=entities_str, text=text)
        swapping_inputs = tokenizer(swapping_prompt, return_tensors="pt").to(device)
        swapping_output = model.generate(
            **swapping_inputs,
            max_new_tokens=max_new_tokens,
            use_cache=True,
            eos_token_id=151645
        )
        
        anonymized_text = tokenizer.decode(swapping_output[0], skip_special_tokens=True)
        anonymized_text = anonymized_text.split("assistant\n", 1)[-1].strip()  # Extract only the assistant's response

        if return_entities:
            return anonymized_text, entities_str

        return anonymized_text

    # Entity Restoration Mode using entity_swapping
    elif mode == "entity_swapping" and entity_mapping:
        # Reverse the entity mapping
        if reverse_mapping:
            reversed_mapping = []
            for line in entity_mapping.splitlines():
                if ':' in line:  # Ensure the line contains a colon
                    left, right = map(str.strip, line.split(":", 1))  # Split and strip spaces
                    reversed_mapping.append(f"{right} : {left}")  # Reverse and format
            entity_mapping = "\n".join(reversed_mapping)
        

        # Create the swapping prompt with the aggregated reversed mappings
        swapping_prompt = formats["entity_swapping"].format(entities=entity_mapping, text=text)
        swapping_inputs = tokenizer(swapping_prompt, return_tensors="pt").to(device)
        swapping_output = model.generate(
            **swapping_inputs,
            max_new_tokens=max_new_tokens,
            use_cache=True,
            eos_token_id=151645
        )

        # Decode and extract the restored text
        output_text = tokenizer.decode(swapping_output[0], skip_special_tokens=True)
        output_text = output_text.split("assistant\n", 1)[-1].strip()  # Extract only the assistant's response

        return output_text

    # Other modes (sensitivity, complexity, entity_detection)
    else:
        prompt = formats[mode].format(text=text)
        model_inputs = tokenizer(prompt, return_tensors="pt").to(device)
        generation_output = model.generate(
            **model_inputs,
            max_new_tokens=5,
            use_cache=True,
            eos_token_id=151645
        )
        full_output = tokenizer.decode(generation_output[0], skip_special_tokens=True)
        
        if mode in ["sensitivity", "complexity"]:
            assistant_text = full_output.split("assistant\n", 1)[-1].strip()
            return assistant_text
        elif mode == "entity_detection":
            return postprocess_entity_recognition(full_output)

# Function to parse entity detection output
def postprocess_entity_recognition(detection_output: str) -> dict:
    output_json = {}
    entity_pattern = re.compile(
        r'(?P<entity>[\w\s]+)--(?P<type>[\w]+)--(?P<random>[\w\s]+)--(?P<generalizations>.+)'
    )
    generalization_pattern = re.compile(r'([\w\s]+)::([\w\s]+)')
    
    lines = detection_output.strip().split("\n")
    for line in lines:
        match = entity_pattern.search(line)
        if match:
            entity_name = match.group("entity").strip()
            entity_type = match.group("type").strip()
            random_replacement = match.group("random").strip()
            
            generalizations = []
            for gen_match in generalization_pattern.findall(match.group("generalizations")):
                first, second = gen_match
        
                # Check if the first part is a digit (score) and swap if needed
                if first.isdigit() and not second.isdigit():
                    score = first
                    label = second
                    generalizations.append([label.strip(), score.strip()])
                    
                elif not first.isdigit() and second.isdigit():
                    label = first
                    score = second
                    generalizations.append([label.strip(), score.strip()])

                
                

            output_json[entity_name] = {
                "TYPE": entity_type,
                "RANDOM": random_replacement,
                "GENERAL": generalizations
            }
    return output_json

# Function to select entities based on config
def select_entities_based_on_json(prediction_json, entity_json):
    entities = []
    for key, value in prediction_json.items():
        entity_type = value["TYPE"]
        if entity_type.upper() in entity_json:
            anonymization_type = entity_json[entity_type]
            if anonymization_type == "RANDOM":
                entities.append([key, value["RANDOM"]])
            elif "GENERAL" in anonymization_type:
                intensity = anonymization_type.split(" ")[1]
                if intensity == "LOW" and value["GENERAL"]:
                    entities.append([key, value["GENERAL"][0][0]])
                elif intensity == "MEDIUM":
                    for gen in value["GENERAL"]:
                        if int(gen[1]) >= 4:
                            entities.append([key, gen[0]])
                            break
                elif intensity == "HIGH":
                    if value["GENERAL"]:
                        entities.append([key, value["GENERAL"][0][0]])
    return entities

We have a community picnic at Greenfield Park, it is on thursday at 11 AM. Write me an e-mail annoucment!

text = "We have a community picnic at Greenfield Park, it is on thursday at 11 AM. Write me an e-mail annoucment!"

# Generate sensitivity score
sensitivity_score = model_inference(text, mode="sensitivity")
print(f"Sensitivity Score: {sensitivity_score}")

# Generate complexity score
complexity_score = model_inference(text, mode="complexity")
print(f"Complexity: {complexity_score}")

Sensitivity Score: 0
Complexity: 3

# Anonymize the text
anonymized_text = model_inference(text, mode="anonymization")
print(f"Anonymized Text: {anonymized_text}\n")

# Restore the original text
anonymized_text, entity_mapping = model_inference(text, mode="anonymization", return_entities=True)
print(f"Entity Mapping:\n{entity_mapping}\n")
restored_text = model_inference(anonymized_text, mode="entity_swapping", entity_mapping=entity_mapping, reverse_mapping=True)
print(f"Restored Text: {restored_text}")

Anonymized Text: We have a community picnic at Sunnyvale Park, it is on A Day of the Week at Morning. Write me an e-mail announcement!

Entity Mapping:
Greenfield Park : Sunnyvale Park
thursday : A Day of the Week
11 AM : Morning

Restored Text: We have a community picnic at Greenfield Park, it is on thursday at 11 AM. Write me an e-mail announcement!