Combining open-source PII redaction with closed-model analysis in healthcare using Llama 3.1, MedSpacy and GPT-4o

import random
import wandb
from datasets import load_dataset
from openai import OpenAI
import json
from typing import Dict

# Constants
FIRST_NAMES = ["John", "Jane", "Michael", "Emily", "Robert", "Jessica", "David", "Sophia", "Matthew", "Amanda"]
LAST_NAMES = ["Smith", "Johnson", "Williams", "Brown", "Jones", "Garcia", "Davis", "Miller", "Wilson", "Martinez"]
DATES = [f"2024-0{i}-15" for i in range(1, 10)]
MEDICAL_RECORD_IDS = [f"MRN-{random.randint(1000, 9999)}" for _ in range(100)]

REPORT_STYLES = [
    "Narrative Style", "SOAP Style", "Checklist Style", "Bullet Point Summary Style",
    "Medico-Legal Style", "Case Summary Style", "Discharge Summary Style", "Teaching Style",
    "Diagnostic-Driven Style", "Interdisciplinary Communication Style", "Layperson Style",
    "Comparative Style", "Emergency Note Style"
]


def generate_prompt(
    section_text: str,
    first_name: str,
    last_name: str,
    date: str,
    record_id: str,
    style: str
) -> str:
    """
    Generate a formatted prompt for clinical report generation.
    """
    style_templates = {
        "Narrative Style": """
Write a detailed clinical report as a narrative. Include all relevant patient details, findings, and recommendations in prose format. If appropriate based on the patient's condition, include any necessary medication prescriptions with dosage and duration.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "SOAP Style": """
Write a clinical report following the SOAP format (Subjective, Objective, Assessment, Plan). Ensure clarity and proper sectioning. In the Plan section, include any appropriate medication prescriptions with specific dosage and duration if warranted by the patient's condition.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Checklist Style": """
Write a clinical report as a checklist. Focus on concise and structured items for each aspect of the case. Include a medications section if any prescriptions are warranted, with clear dosage and duration details.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Bullet Point Summary Style": """
Write a clinical report as a bullet point summary. Highlight key findings, diagnosis, and next steps. Include any necessary medication prescriptions as separate bullet points with clear dosage and duration information.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Medico-Legal Style": """
Write a formal medico-legal report. Include all findings with precise language suitable for legal review. Document any prescribed medications with complete details including dosage, duration, and clinical rationale.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Case Summary Style": """
Summarize the case in a concise format suitable for a referral or handoff. Include details of any prescribed medications, including dosage and duration, if warranted by the patient's condition.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Discharge Summary Style": """
Write a discharge summary focusing on follow-up care and post-visit recommendations. Include a detailed medication section with any new prescriptions, including dosage, duration, and instructions for use.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Teaching Style": """
Write a clinical report suitable for medical training purposes. Include reasoning and explanations for the findings. If prescribing any medications, explain the clinical rationale for each prescription including choice of drug, dosage, and duration.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Diagnostic-Driven Style": """
Write a report focused on the diagnostic process. Highlight the reasoning behind the diagnosis. Include any therapeutic decisions, including medication prescriptions with dosage and duration when appropriate.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Interdisciplinary Communication Style": """
Write a clinical report intended for communication with other healthcare professionals. Detail any medication changes or new prescriptions, including full prescribing information and rationale.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Layperson Style": """
Write a clinical report in simple language that a patient without medical knowledge can understand. If any medications are prescribed, explain them in plain language, including why they're needed and how to take them properly.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Comparative Style": """
Write a clinical report comparing the current findings with previous evaluations. Include any changes in medication regimen, new prescriptions, or discontinued medications with complete details.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
""",
        "Emergency Note Style": """
Write a brief and urgent clinical report as if documenting an emergency case. Include any immediate medication interventions or prescriptions with precise dosing instructions.

Patient Name: {first_name} {last_name}
Date: {date}
Medical Record Number: {record_id}

### Input:
{section_text}

### Synthetic Report:
"""
    }

    if style not in style_templates:
        raise ValueError(f"Unknown style: {style}")

    return style_templates[style].format(
        first_name=first_name,
        last_name=last_name,
        date=date,
        record_id=record_id,
        section_text=section_text
    )
def generate_random_identifiers():
    """Generate a set of random identifiers for a patient."""
    return (
        random.choice(FIRST_NAMES),
        random.choice(LAST_NAMES),
        random.choice(DATES),
        random.choice(MEDICAL_RECORD_IDS)
    )
def generate_gpt_synthetic_report(section_text: str, client: OpenAI) -> str:
    """Generate a synthetic medical report using the OpenAI API."""
    first_name, last_name, date, record_id = generate_random_identifiers()
    style = random.choice(REPORT_STYLES)
    
    prompt = generate_prompt(
        section_text=section_text,
        first_name=first_name,
        last_name=last_name,
        date=date,
        record_id=record_id,
        style=style
    )

    response = client.chat.completions.create(
        model="gpt-4o-2024-08-06",
        temperature=0.0,
        messages=[
            {"role": "system", "content": f"You are a medical assistant creating clinical reports in {style}."},
            {"role": "user", "content": prompt}
        ]
    )

    return response.choices[0].message.content.strip()

def main():
    # Initialize OpenAI client
    OPENAI_API_KEY = ""

    client = OpenAI(api_key=OPENAI_API_KEY)
    
    # Initialize WandB
    wandb.init(project="synthetic-clinical-reports")
    
    # Load dataset
    dataset = load_dataset("har1/MTS_Dialogue-Clinical_Note")
    first_hundred = dataset["train"].select(range(100))
    
    # Generate reports
    synthetic_reports = []
    for row in first_hundred:
        section_text = row["section_text"]
        synthetic_report = generate_gpt_synthetic_report(section_text, client)
        synthetic_reports.append({
            "id": row["ID"],
            "original_section_text": section_text,
            "synthetic_report": synthetic_report
        })
    
    # Log to WandB and save
    table = wandb.Table(columns=["ID", "Original Section Text", "Synthetic Report"])
    for report in synthetic_reports:
        table.add_data(
            report["id"], 
            report["original_section_text"], 
            report["synthetic_report"]
        )
    wandb.log({"Synthetic Reports": table})
    
    # Save to file
    with open("./synthetic_reports.json", "w") as f:
        json.dump(synthetic_reports, f, indent=4)


if __name__ == "__main__":
    main()

import transformers
import torch
import json
import weave; weave.init("masking_pii")


# Load the dataset from a local file
with open("synthetic_reports.json", "r") as f:
    dataset = json.load(f)


# Load the Llama model using the transformers pipeline
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
)


# List of HIPAA identifiers to mask
PII_CATEGORIES = [
    "Names", "Geographic information (addresses, cities, counties, zip codes)", 
    "Dates (birth, admission, discharge, death, except year)", "Phone Numbers", 
    "Fax Numbers", "Email Addresses", "Social Security Numbers", 
    "Medical Record Numbers", "Health Plan Beneficiary Numbers", "Account Numbers", 
    "Certificate/License Numbers", "Vehicle Identifiers (including license plates)", 
    "Device Identifiers", "Web URLs", "IP Addresses", "Biometric Identifiers (e.g., fingerprints)", 
    "Full-Face Photographic Images", "Any other unique identifier or code"
]


def postprocess_inputs(inputs: dict) -> dict:
    # Exclude the 'text' input to avoid logging sensitive data
    return {k: v for k, v in inputs.items() if k != 'text'}


@weave.op(postprocess_inputs=postprocess_inputs)
def mask_pii(text):
    """
    Use the Llama model to mask PII in the input text.
    It replaces sensitive information with '[redacted]'.
    """
    # Format the messages for the new prompt style
    messages = [
        {"role": "system", "content": "You are an advanced language model tasked with identifying and masking personally identifiable information (PII) in text."},
        {"role": "user", "content": f"""
The following text contains personally identifiable information (PII). 
Please identify all instances of the following PII categories: {', '.join(PII_CATEGORIES)} 
and replace them with '[redacted]'. Ensure that the rest of the text remains unchanged. DO NOT RESPOND WITH ANY PII INFORMATION WHATSOEVER!


Text: "{text}"

Masked Text:
        """}
    ]
    
    # Generate masked text
    response = pipeline(messages, max_new_tokens=2048)
    print(response[0])
    return response[0]["generated_text"][2]['content']



# Iterate over the dataset and mask PII
masked_data = []
for record in dataset:
    original_text = record["synthetic_report"]
    response = mask_pii(original_text)
    
    # Assuming the pipeline response is a list of dictionaries
    masked_text = response  # Extract the generated text directly
    
    masked_data.append({
        "id": record["id"],
        "original_text": original_text,
        "masked_text": masked_text.strip()  # Ensure whitespace is trimmed
    })
    print(masked_text)  # Print the masked text for debugging


# Save the masked dataset to a new JSON file
with open("masked_ds.json", "w") as f:
    json.dump(masked_data, f, indent=4)


print("PII masking completed. Masked dataset saved to 'masked_ds.json'.")



def postprocess_inputs(inputs: dict) -> dict:
# Exclude the 'text' input to avoid logging sensitive data
return {k: v for k, v in inputs.items() if k != 'text'}

@weave.op(postprocess_inputs=postprocess_inputs)
def mask_pii(text):
	# rest of the function 

import json
from openai import OpenAI
import weave; weave.init("masking_pii_double_check")


# Set up OpenAI API key
OPENAI_API_KEY = ""
oclient = OpenAI(api_key=OPENAI_API_KEY)

def analyze_report(text):
    """
    Use GPT-4o to analyze a masked clinical report for possible medication interactions
    or potential misdiagnoses.
    """
    messages = [
        {"role": "system", "content": "You are a medical expert analyzing clinical reports for medication interactions and potential misdiagnoses."},
        {"role": "user", "content": f"""
The following is a masked clinical report. Please analyze it for:
1. Any potential medication interactions based on the prescribed medications.
2. Possible misdiagnoses based on the symptoms and other clinical details provided.

Masked Clinical Report:
{text}

Analysis:
        """}
    ]
    
    # Generate analysis
    response = oclient.chat.completions.create(
        model="gpt-4o-2024-08-06",
        temperature=0.5,
        messages=messages
    )
    return response.choices[0].message.content.strip()

def main():
    # Load the masked dataset
    with open("masked_ds.json", "r") as f:
        masked_data = json.load(f)

    analysis_results = []

    # Analyze each masked report
    for record in masked_data:
        original_text = record["masked_text"]
        analysis = analyze_report(original_text)
        analysis_results.append({
            "id": record["id"],
            "masked_text": original_text,
            "analysis": analysis
        })
        print(f"Analysis for ID {record['id']} completed.")  # Debugging output

    # Save the analysis results to a new JSON file
    output_path = "analysis_results.json"
    with open(output_path, "w") as f:
        json.dump(analysis_results, f, indent=4)

    print(f"Analysis completed. Results saved to '{output_path}'.")

if __name__ == "__main__":
    main()

import spacy
import medspacy
import json

# Load MedSpaCy with the med7 pipeline
med7 = spacy.load("en_core_med7_lg")
nlp = medspacy.load(med7)

# Load the masked dataset
with open("masked_ds.json", "r") as f:
    masked_dataset = json.load(f)

# Process the masked data
processed_data = []

for record in masked_dataset:
    # Read the masked text
    masked_text = record["masked_text"]
    
    # Use medspacy to process the text
    doc = nlp(masked_text)
    
    # Extract entities
    entities = [{"text": ent.text, "label": ent.label_} for ent in doc.ents]
    
    # Append results
    processed_data.append({
        "id": record["id"],
        "masked_text": masked_text,
        "entities": entities
    })

# Write the processed data to a new JSON file
with open("processed_masked_ds.json", "w") as f:
    json.dump(processed_data, f, indent=4)

print("Processing completed. Results saved to 'processed_masked_ds.json'.")