import gradio as gr
from transformers import pipeline

# Load a lightweight model
generator = pipeline("text2text-generation", model="facebook/bart-large", tokenizer="facebook/bart-large")

# Function to construct the anonymization prompt
def construct_prompt(input_text):
    prompt = f"""
    Please rewrite the input text by writting the text in double.
    
    "{input_text}"
    """
    return prompt.strip()

# Function to process the input text and generate the anonymized output
def anonymize_text(input_text):
    # Construct the instruction prompt
    prompt = construct_prompt(input_text)
    
    # Generate the output using the model
    response = generator(prompt, max_length=512, num_return_sequences=1)
    generated_text = response[0]['generated_text']
    
    # Attempt to parse the JSON output
    try:
        result = eval(generated_text)  # Convert the output string to a Python dictionary
        anonymized_text = result.get("anonymized_text", input_text)
        name_mapping = result.get("name_mapping", {})
    except Exception as e:
        anonymized_text = input_text
        name_mapping = {}
    
    return anonymized_text, name_mapping

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Text Anonymizer")
    gr.Markdown("Enter text containing personal names, and the model will anonymize it by replacing the names with pseudonyms. The app will also return a dictionary linking original names to their replacements.")
    
    input_text = gr.Textbox(label="Input Text", placeholder="Enter text here...")
    anonymized_text = gr.Textbox(label="Anonymized Text", interactive=False)
    name_mapping = gr.JSON(label="Name Mapping")
    
    def process_text(input_text):
        return anonymize_text(input_text)
    
    submit_button = gr.Button("Anonymize")
    submit_button.click(process_text, inputs=[input_text], outputs=[anonymized_text, name_mapping])

# Launch the app
demo.launch()