import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
import gradio as gr

# Load the tokenizer and model
model_name = "iiiorg/piiranha-v1-detect-personal-information"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)

# Set device to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def apply_redaction(masked_text, start, end, pii_type, aggregate_redaction):
    for j in range(start, end):
        masked_text[j] = ''
    if aggregate_redaction:
        masked_text[start] = '[redacted]'
    else:
        masked_text[start] = f'[{pii_type}]'

def mask_pii(text, aggregate_redaction=True):
    # Tokenize input text
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    # Get the model predictions
    with torch.no_grad():
        outputs = model(**inputs)

    # Get the predicted labels
    predictions = torch.argmax(outputs.logits, dim=-1)

    # Convert token predictions to word predictions
    encoded_inputs = tokenizer.encode_plus(text, return_offsets_mapping=True, add_special_tokens=True)
    offset_mapping = encoded_inputs['offset_mapping']

    masked_text = list(text)
    is_redacting = False
    redaction_start = 0
    current_pii_type = ''

    for i, (start, end) in enumerate(offset_mapping):
        if start == end:  # Special token
            continue

        label = predictions[0][i].item()
        if label != model.config.label2id['O']:  # Non-O label
            pii_type = model.config.id2label[label]
            if not is_redacting:
                is_redacting = True
                redaction_start = start
                current_pii_type = pii_type
            elif not aggregate_redaction and pii_type != current_pii_type:
                # End current redaction and start a new one
                apply_redaction(masked_text, redaction_start, start, current_pii_type, aggregate_redaction)
                redaction_start = start
                current_pii_type = pii_type
        else:
            if is_redacting:
                apply_redaction(masked_text, redaction_start, end, current_pii_type, aggregate_redaction)
                is_redacting = False

    # Handle case where PII is at the end of the text
    if is_redacting:
        apply_redaction(masked_text, redaction_start, len(masked_text), current_pii_type, aggregate_redaction)

    return ''.join(masked_text)

# Define the function for Gradio interface
def redact_text(text, aggregate_redaction):
    return mask_pii(text, aggregate_redaction)

# Create Gradio Interface
demo = gr.Interface(
    fn=redact_text,
    inputs=[
        gr.Textbox(lines=5, label="Enter Text with Potential PII"),
        gr.Checkbox(label="Aggregate Redaction", value=True)
    ],
    outputs="text",
    title="PII Detection and Redaction",
    description="This application detects personal identifiable information (PII) and redacts it from the provided text. You can choose to either aggregate all PII redaction into a single '[redacted]' label or keep each PII type labeled individually.",
    examples=[
        ["John Doe's phone number is 123-456-7890, and his email is john.doe@example.com."],
        ["Jane was born on 12th August, 1990 and her SSN is 987-65-4321."]
    ]
)

if __name__ == "__main__":
    demo.launch()