전자 건강 기록(EHR)을 이용한 역학 질의응답을 위한 검색 증강(Retrieval augmented) text-to-SQL 생성

이 저장소는 다음 논문의 코드를 포함하고 있습니다:

@misc{ziletti2024retrieval,
title={Retrieval augmented text-to-SQL generation for epidemiological question answering using electronic health records},
author={Angelo Ziletti and Leonardo D'Ambrosi},
...

Preprint(사전 출판본)는 여기에서 확인하실 수 있습니다: Ziletti and D'Ambrosi, https://arxiv.org/abs/2403.09226

이 논문은 NAACL 2024 Clinical NLP Workshop (https://clinical-nlp.github.io/2024/)에 채택되었습니다.

귀하의 작업이나 연구에 이 코드를 사용하는 경우 이 저작물을 인용해 주세요.

다음은 절차의 요약된 워크플로우입니다:

더 자세한 내용은 원고를 참조하십시오: https://arxiv.org/abs/2403.09226

라벨링된 질문-SQL 쌍은 dataset 폴더 내의 text2sql_epi_dataset_omop.xlsx 파일에 있습니다.

저희는 실험에 Python 3.11을 사용했습니다.

GitHub의 master 브랜치에서 최신 버전을 다음과 같이 설치하십시오:

git clone <GITHUB-URL>
cd text-to-sql-epi-ehr-naacl2024
pip install -r requirements.txt

먼저, dataset 폴더에 제공된 Excel 파일로부터 Pickle 파일(쿼리 라이브러리)을 생성해야 합니다.

cd scripts
python run_querylib_calc.py

이 Pickle 파일은 이후 쿼리 생성 단계에서 검색 증강 생성 (RAG, Retrieval Augmented Generation)을 수행하는 데 사용됩니다.

예측을 수행하려면 질문을 큰따옴표로 묶어 prediction_pipeline.py 스크립트를 실행하십시오. 예시:

cd scripts
python prediction_pipeline.py --question "How many women with atopic dermatitis?"

결과는 다음과 같이 반환되어야 합니다:

Question: How many women with atopic dermatitis?
SQL template:
SELECT COUNT(DISTINCT p.person_id) AS female_patients_with_atopic_dermatitis
...

쿼리를 실행하려면 다음 Google Cloud의 데이터셋 링크를 사용하세요: https://console.cloud.google.com/marketplace/product/hhs/synpuf 링크를 열고(Google 계정 필요), 귀하의 Google 계정에서 쿼리를 실행할 수 있습니다.

Google BigQuery에서 쿼리를 실행 가능하게 만들려면, 생성된 테이블 이름을 Google Query Table Names 규격에 맞게 수정해야 합니다. 이는 단순히 테이블의 이름을 변경하는 작업입니다.

예를 들어,

SELECT COUNT(DISTINCT p.person_id) AS female_patients_with_atopic_dermatitis
FROM person AS p
JOIN condition_occurrence AS co ON p.person_id = co.person_id
...

위의 쿼리는 위의 BigQuery 인스턴스에서 실행하기 위해 다음과 같이 변경되어야 합니다.

SELECT
COUNT(DISTINCT p.person_id) AS female_patients_with_atopic_dermatitis
FROM `bigquery-public-data.cms_synthetic_patient_data_omop.person` AS p
...

만약 자체적인 Snowflake 데이터 웨어하우스 (Data Warehouse)를 보유하고 있다면, 이러한 변경을 수행할 필요가 없을 수도 있습니다.

저장소(Repository)에서는 위 그림에서 보여주는 것처럼 의료 코딩 (medical coding)이 프로세스에 어떻게 통합될 수 있는지 보여주기 위해 의료 코딩용 모의 버전 (mockup version)을 제공합니다. 코딩은 SNOMED 온톨로지 (ontology) 내에서 수행되는데, 이는 OMOP 공통 데이터 모델 (OMOP-CDM)에 저장된 데이터의 기반이 되는 온톨로지이기 때문입니다. 이 테이블은 일반적으로 OMOP-CDM에서 CONCEPT_TABLE이라고 불립니다.

우리는 dataset 폴더에 medcodes_mockup.xlsx라는 이름의 작은 모의 테이블을 준비했습니다. 여기에는 OMOP-CDM의 데이터 구조를 따르는 합성 데이터 (synthetic data)가 포함되어 있습니다.
전체 온톨로지는 이곳 국립 의학 도서관 (National Library of Medicine)에서 다운로드할 수 있습니다.
(이용 약관에 동의해야 함을 유의하세요)

SNOMED-CT에 대한 더 자세한 정보는 예를 들어 이곳에서 확인할 수 있습니다.

먼저, 제공된 Excel 파일로부터 Pickle 파일 (medcodeonto 라이브러리)을 생성해야 합니다.

cd scripts
python run_medcoding_calc.py

이 Pickle 파일은 이후 Pickle 파일 내에서 의료 코드를 검색하는 데 사용됩니다. 파일을 실행하면 코딩이 어떻게 수행되는지에 대한 예시도 함께 제공됩니다.

이것은 단지 목업 (mockup) 구현이라는 점에 유의하십시오. 프로덕션 환경에 적합한 버전을 위해서는, 예를 들어 Snowflake 데이터베이스의 전체 CONCEPT 테이블을 인덱싱하여 벡터 데이터베이스 (vector database, 예: qdrant)를 사용할 것을 권장합니다.

Medical coding compilation 단계를 실행한 후에는, 전체 파이프라인인 SQL 쿼리 생성 (SQL query generation) + 의료 코딩 (medical coding)을 실행할 준비가 됩니다.

다음은 예시입니다:

cd scripts
python prediction_pipeline.py --med_coding True --question "How many females with atopic dermatitis"

실행 결과는 다음과 같습니다:

Use medical coding: True
Use Snowflake database: False
No sentence-transformers model found with name BAAI/bge-large-en-v1.5. Creating a new one with MEAN pooling.
...

의료 코딩 단계는 여기에서 예시용으로만 제공된다는 점을 유의하십시오. 신뢰할 수 있는 결과를 얻으려면 전체 SNOMED 온톨로지 (ontology)에 대한 접근 권한이 필요합니다 (Medical coding compilation 섹션 참조).