prompt.py

# ----------- Task 1: Prompt engineering: reorganize the X-Ray report findings into predefined anatomical regions

# import packages
import openai
from openai import OpenAI
import json
import os
from tqdm import tqdm
import argparse


# Set up argument parsing
parser = argparse.ArgumentParser()
parser.add_argument('-i', "--input_path", type=str, help='Path to the input JSON file containing reports')
parser.add_argument('-o', "--output_path", type=str, help='Path to the output JSON file to save categorized reports')
parser.add_argument('-k', "--openAI_API_key", type=str, required=True, help='Your OpenAI API key')
args = parser.parse_args()


# Set the input and output file paths 
if args.input_path:
   input_file_path = args.input_path
else:
   input_file_path = './data/annotation_quiz_all.json'

if args.output_path:
   output_file_path = args.output_path
else:
   output_file_path = './data/annotation.json' # Update existing file with categorized reports in val set


# Set your API key for OpenAI
client = OpenAI(api_key= args.openAI_API_key) # Specify your OpenAI API key here


# Function to prompt gpt-4o-mini to categorize the findings 
def categorize_findings(report):
    # Create a chat completion request using a structured prompt
    chat_completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": """Categorize the findings of a chest X-ray report into predefined anatomical regions: bone, heart, lung, and mediastinal. 
                    If a finding does not clearly belong to these categories, classify it under 'others'. Read each sentence carefully. Determine the main anatomical focus of each sentence:
                    - If a sentence discusses any findings related to bones, categorize it under 'bone'.
                    - If it pertains to the heart, categorize it under 'heart'. 
                    - If a sentence discusses any findings related to the lungs or associated structures, categorize it under 'lung'.
                    - If it mentions any findings related to the mediastinal area, categorize it under 'mediastinal'.
                    - If a sentence does not fit any of the above categories or is ambiguous, place it under 'others'.
                    Provide the output as a JSON object with each category listing relevant sentences from the report in **plain text** without extra double quotes around the sentences. 
                    The format should be: {"bone": "", "heart": "", "lung": "", "mediastinal": "", "others": ""}.
                    """
            },
            {
                "role": "user",
                "content": report
            }
        ],
        response_format= {"type": "json_object"}
    )
    # Extract and return the model's response
    return chat_completion.choices[0].message.content


# Test code with 1 report
sample_report = "The cardiomediastinal silhouette and pulmonary vasculature are within normal limits in size. The lungs are mildly hypoinflated but grossly clear of focal airspace disease, pneumothorax, or pleural effusion. There are mild degenerative endplate changes in the thoracic spine. There are no acute bony findings."
categorized_report = categorize_findings(sample_report)
print(categorized_report)

result = json.loads(categorized_report)
#print(result)
#print(result['lung'])


# -----------  For all reports

# Get all reports

# Read the JSON file
with open(input_file_path, 'r') as file:
    data = json.load(file)
    val_reports = data.get('val', []) # retrieve the value associated with the key 'val' from the dictionary

print("Num of Reports in Val set: ",len(val_reports))



# Categorize findings for all reports with batching

# Batch size
batch_size = 10  # set batch size
categorized_results = []

# Process the reports in batches
for i in tqdm(range(0, len(val_reports), batch_size), desc="Processing reports"):
    batch = val_reports[i:i + batch_size]  # Get the current batch of reports
    
    for report in batch:
        try:
            result = categorize_findings(report['original_report'])  # Get output from the model
            result = json.loads(result)  # Load string from JSON object
        except Exception as e:
            print(f"Error processing report {report['id']}: {e}")
            continue

        dict_results = {'id': report['id'], 'report': result, 'split': report['split']}
        categorized_results.append(dict_results)

    # Also replace the original file with the updated results
    with open(input_file_path, 'r') as file: # Read the JSON file
        data = json.load(file)
    data['val'] = categorized_results # Update the 'val' key with the categorized results

    with open(output_file_path, 'w') as file: # Write the updated JSON back to a new file
        json.dump(data, file, indent=4)
    
    print("File updated successfully.")