README

Overview

This repository contains the code necessary to replicate some of the results of our research paper on trial prescreening using LLMs. The code converts the XML files from Harvard's N2C2 dataset into a set of images to replicate a real world situation, vectorizes them using VoyageAI's multimodal embedding model and finally uses OpenAI's o1 reasoning model to make an assessment given a criterion and the top K most relevant images (with retrieval guidelines or not).

Installation

1. Clone this repository
2. Navigate to the cloned directory
3. Create and activate a virtual environment:

python -m venv .venv 
source .venv/bin/activate

4. Install dependencies:

pip install -r requirements.txt

5. Copy the file .env.template to a new .env file and fill it with an OpenAI and a VoyageAI API Key and optionally with Langfuse API Keys for tracing.

6. Unzip XML files from the N2C2 dataset in ./data. Note that it already contains a fake sample that you may use to run this code.

Usage

Below are examples of how to run the four core scripts. Adjust paths and parameters as needed.

1. scripts/generate_dataset.py

Processes XML files from the n2c2 Harvard dataset as follows:

1. Converts the text medical record of each patient in a set of images of standard page size
2. Builds a dataframe containing 1 row for each pair Patient x Criterion and containing the expected output (met/not met) as well as paths to the relevant images for that patient

Example command:

python -m scripts.generate_dataset <OUTPUT_FILE_NAME> 

With:

• OUTPUT_FILE_NAME: a file name without extension. It will be used to create 2 files in the ./outputs directory : <OUTPUT_FILE_NAME>.csv and <OUTPUT_FILE_NAME>.pkl.

Images of medical records will be created in ./outputs/images. It takes ~1 min to generate images for the 200 patients of the dataset.

2. scripts/generate_embeddings.py

Vectorizes the images created at the previous step using VoyageAI's multimodal embedding model and saves the embedding in a FAISS vector database.

Example command:

python -m scripts.generate_embeddings <INPUT_FILE_NAME> 

With:

• INPUT_FILE_NAME: the same file name that you chose as an output for the previous step

Images of medical records will be created in ./outputs/images. It takes ~30 min to generate embbedings of all images.

3. scripts/assess_criteria.py

Uses the dataframe and images created at the previous step to run eligibility assessments using OpenAI's o1 model for every patient x criterion pair.

Example command:

python -m scripts.assess_criteria <INPUT_FILE_NAME> <OUTPUT_PATH> --test_size <TEST_SIZE> --run_on_test_split --use_refined_criteria --top_k_images <K> --use_retrieval_guidelines --<custom_trace_keyword>=<value>

With:

• INPUT_FILE_NAME: the same file name that you chose as an input for step 2
• OUTPUT_PATH: a directory relative to ./outputs where results will be saved. If it does not exist, it will be created
• TEST_SIZE: the ratio of patient to put in the test set (default is 0.9, to use 90% of patients for testing and 10% for prompt engineering / optimizations).
• --run_on_test_split: if present, the pipeline will be executed on the test split, if not, the pipeline will be executed on the train split.
• --use_refined_criteria: if present, the refined versions of criteria from the Zero-Shot Clinical Trial Patient Matching with LLMs paper will be used. If not, the raw criteria from the Harvard's n2c2 dataset will be used.
• --top_k_images=<K>: the number of most relevant images to retrieve for each criterion
• --use_retrieval_guidelines: if present, a set of retrieval guidelines will be generated for each criterion and used to retrieve the K most relevant images for each guideline. As a result, the criterion assessment will likely be done with more than K images.
• --<custom_trace_keyword>=<value>: any number of key value pairs to be added as metadata on Langfuse traces.

To run the pipeline on a small set of examples, use :

python -m scripts.assess_criteria <INPUT_FILE_NAME> <OUTPUT_PATH> --test_size 0.995

By default, 100 queries will be sent to OpenAI in parallel, with a maximum of 500 queries per minute. You can control this with below parameters in your .env file:

• OPENAI_MAX_REQUESTS_PER_MINUTE=500
• MAX_OPENAI_CONCURRENCY=100

4. Results

There are 2 ways to generate metrics and plots. The first one uses the results.csv files created at the previous step. And the second one uses Langfuse Traces. The second method is preferred because it provides an aggregated view as well as latency and costs metrics.

4.1. scripts/generate_metrics.py

Uses the assessments generated at the previous step to generate various metrics

Example command:

python -m scripts.generate_metrics <INPUT_PATH> 

With:

• INPUT_PATH: a directory relative to ./outputs where results were saved at the previous step. Files containing the metrics will be saved in the same directory.

4.2. Metrics using Langfuse traces

If LangFuse API Key was added at step 3., traces were uploaded to Langfuse and we can use them to analyze to results. To do so, export Langfuse traces to ./outputs/results_aggregated/all_langfuse_traces.csv and run below scripts.

4.2.1. scripts/process_langfuse_traces.py

Uses the langfuse traces and the ground truth file to generate a file containing all assessments.

Example command:

python -m scripts.process_langfuse_traces <ground_truth_input_path> <langfuse_traces_input_path>

With:

• ground_truth_input_path: the directory relative to ./outputs where ground truth file was saved at step 1.
• langfuse_traces_input_path: the name of file containing langfuse traces, e.g. all_langfuse_traces.csv

4.2.2. scripts/generate_metrics_from_traces.py

Uses the processed langfuse traces generated at previous step to generate plots and metrics.

Example command:

python -m scripts.generate_metrics_from_traces

Metrics and plots will be saved in .outputs/results_aggregated