| import os |
| import csv |
| import json |
| import re |
| import random |
| from collections import defaultdict |
| from typing import Dict, List, Set, Tuple, Optional |
| from tqdm import tqdm |
|
|
|
|
| SNOMED_TAG_REGEX = re.compile(r"\s*\([^()]*\)$") |
|
|
|
|
| def clean_text(text: str) -> str: |
| """Normalize text: strip surrounding whitespace, quotes, and collapse multiple spaces.""" |
| if not text: |
| return "" |
| text = text.strip().strip('"').strip("'") |
| text = re.sub(r"\s+", " ", text) |
| return text |
|
|
|
|
| def strip_snomed_tag(term: str) -> Tuple[str, Optional[str]]: |
| """ |
| Strips trailing semantic tags from SNOMED terms (e.g. 'Asthma (disorder)' -> 'Asthma', '(disorder)'). |
| Returns (cleaned_term, tag_or_none). |
| """ |
| match = SNOMED_TAG_REGEX.search(term) |
| if match: |
| tag = match.group(0).strip() |
| cleaned = term[:match.start()].strip() |
| if len(cleaned) > 1: |
| return cleaned, tag |
| return term, None |
|
|
|
|
| class AthenaVocabularyParser: |
| """ |
| High-performance parser for Athena OMOP CDM exports: |
| - CONCEPT.csv |
| - CONCEPT_SYNONYM.csv |
| Extracts and aligns SNOMED CT, LOINC, and RxNorm concepts and synonyms. |
| """ |
|
|
| def __init__( |
| self, |
| concept_file: str = "vocabularies/CONCEPT.csv", |
| synonym_file: str = "vocabularies/CONCEPT_SYNONYM.csv", |
| target_vocabs: Optional[List[str]] = None, |
| target_languages: Optional[List[str]] = None, |
| include_all_languages: bool = False, |
| strip_snomed_tags: bool = True, |
| min_synonyms: int = 2, |
| max_synonyms: int = 25, |
| ): |
| self.concept_file = concept_file |
| self.synonym_file = synonym_file |
| self.target_vocabs = set(target_vocabs or ["SNOMED", "LOINC", "RxNorm"]) |
| self.target_languages = set(target_languages or ["4180186"]) |
| self.include_all_languages = include_all_languages |
| self.strip_snomed_tags = strip_snomed_tags |
| self.min_synonyms = min_synonyms |
| self.max_synonyms = max_synonyms |
|
|
| def parse( |
| self, |
| output_dir: str = "data/processed", |
| subset_size: Optional[int] = None, |
| train_ratio: float = 0.8, |
| val_ratio: float = 0.1, |
| test_ratio: float = 0.1, |
| seed: int = 42, |
| ) -> Dict[str, int]: |
| """ |
| Parses CONCEPT.csv and CONCEPT_SYNONYM.csv, cleans synonyms, and splits into train/val/test. |
| """ |
| os.makedirs(output_dir, exist_ok=True) |
| random.seed(seed) |
|
|
| print(f"Step 1: Reading concepts from {self.concept_file}...") |
| valid_concepts: Dict[str, Dict] = {} |
| concept_synonyms: Dict[str, Set[str]] = defaultdict(set) |
|
|
| total_concept_rows = 0 |
| with open(self.concept_file, "r", encoding="utf-8", errors="replace") as f: |
| reader = csv.DictReader(f, delimiter="\t") |
| for row in tqdm(reader, desc="Scanning CONCEPT.csv"): |
| total_concept_rows += 1 |
| vocab = row.get("vocabulary_id", "") |
| invalid_reason = row.get("invalid_reason", "") |
|
|
| |
| if vocab in self.target_vocabs and not invalid_reason: |
| cid = row["concept_id"] |
| cname = clean_text(row.get("concept_name", "")) |
| if not cname: |
| continue |
|
|
| valid_concepts[cid] = { |
| "concept_id": cid, |
| "concept_name": cname, |
| "vocabulary_id": vocab, |
| "concept_code": row.get("concept_code", ""), |
| "concept_class_id": row.get("concept_class_id", ""), |
| "standard_concept": row.get("standard_concept", ""), |
| "domain_id": row.get("domain_id", ""), |
| } |
|
|
| |
| concept_synonyms[cid].add(cname) |
|
|
| |
| if vocab == "SNOMED" and self.strip_snomed_tags: |
| stripped, _ = strip_snomed_tag(cname) |
| if stripped and stripped != cname: |
| concept_synonyms[cid].add(stripped) |
|
|
| print(f"Total valid concepts loaded for {self.target_vocabs}: {len(valid_concepts):,}") |
|
|
| print(f"Step 2: Reading synonyms from {self.synonym_file}...") |
| synonym_rows_processed = 0 |
| with open(self.synonym_file, "r", encoding="utf-8", errors="replace") as f: |
| reader = csv.DictReader(f, delimiter="\t") |
| for row in tqdm(reader, desc="Scanning CONCEPT_SYNONYM.csv"): |
| cid = row.get("concept_id", "") |
| if cid not in valid_concepts: |
| continue |
|
|
| lang = row.get("language_concept_id", "") |
| if not self.include_all_languages and lang not in self.target_languages: |
| continue |
|
|
| sname = clean_text(row.get("concept_synonym_name", "")) |
| if len(sname) < 2 or len(sname) > 300: |
| continue |
|
|
| concept_synonyms[cid].add(sname) |
| synonym_rows_processed += 1 |
|
|
| |
| if valid_concepts[cid]["vocabulary_id"] == "SNOMED" and self.strip_snomed_tags: |
| stripped, _ = strip_snomed_tag(sname) |
| if stripped and stripped != sname: |
| concept_synonyms[cid].add(stripped) |
|
|
| print(f"Synonym rows mapped: {synonym_rows_processed:,}") |
|
|
| |
| trainable_concept_ids = [ |
| cid for cid, syns in concept_synonyms.items() if len(syns) >= self.min_synonyms |
| ] |
| print(f"Concepts with >= {self.min_synonyms} synonyms: {len(trainable_concept_ids):,}") |
|
|
| |
| if subset_size and subset_size < len(trainable_concept_ids): |
| random.shuffle(trainable_concept_ids) |
| trainable_concept_ids = trainable_concept_ids[:subset_size] |
| print(f"Sampled subset of {len(trainable_concept_ids):,} concepts for training.") |
|
|
| |
| trainable_concept_ids.sort() |
| random.shuffle(trainable_concept_ids) |
|
|
| |
| n_total = len(trainable_concept_ids) |
| n_train = int(n_total * train_ratio) |
| n_val = int(n_total * val_ratio) |
|
|
| train_ids = set(trainable_concept_ids[:n_train]) |
| val_ids = set(trainable_concept_ids[n_train : n_train + n_val]) |
| test_ids = set(trainable_concept_ids[n_train + n_val :]) |
|
|
| print(f"Splits -> Train: {len(train_ids):,}, Val: {len(val_ids):,}, Test: {len(test_ids):,}") |
|
|
| |
| split_files = { |
| "train": os.path.join(output_dir, "train.jsonl"), |
| "val": os.path.join(output_dir, "val.jsonl"), |
| "test": os.path.join(output_dir, "test.jsonl"), |
| } |
| split_id_map = {"train": train_ids, "val": val_ids, "test": test_ids} |
|
|
| counts = {} |
| for split_name, c_ids in split_id_map.items(): |
| filepath = split_files[split_name] |
| total_pairs = 0 |
| with open(filepath, "w", encoding="utf-8") as out_f: |
| for cid in c_ids: |
| syn_list = list(concept_synonyms[cid]) |
| if len(syn_list) > self.max_synonyms: |
| |
| cname = valid_concepts[cid]["concept_name"] |
| other_syns = [s for s in syn_list if s != cname] |
| random.shuffle(other_syns) |
| syn_list = [cname] + other_syns[: self.max_synonyms - 1] |
|
|
| record = { |
| "concept_id": cid, |
| "vocabulary_id": valid_concepts[cid]["vocabulary_id"], |
| "concept_code": valid_concepts[cid]["concept_code"], |
| "concept_name": valid_concepts[cid]["concept_name"], |
| "synonyms": syn_list, |
| } |
| out_f.write(json.dumps(record, ensure_ascii=False) + "\n") |
| total_pairs += len(syn_list) |
| counts[split_name] = len(c_ids) |
| counts[f"{split_name}_terms"] = total_pairs |
| print(f"Saved {split_name} -> {filepath} ({len(c_ids):,} concepts, {total_pairs:,} terms)") |
|
|
| |
| meta_filepath = os.path.join(output_dir, "concepts_meta.json") |
| with open(meta_filepath, "w", encoding="utf-8") as f: |
| json.dump(valid_concepts, f, indent=2, ensure_ascii=False) |
| print(f"Saved concept metadata lookup -> {meta_filepath}") |
|
|
| |
| stats_file = os.path.join(output_dir, "dataset_stats.json") |
| stats = { |
| "total_concepts_in_vocab": len(valid_concepts), |
| "trainable_concepts": len(trainable_concept_ids), |
| "splits": counts, |
| "vocabularies": list(self.target_vocabs), |
| } |
| with open(stats_file, "w", encoding="utf-8") as f: |
| json.dump(stats, f, indent=2) |
|
|
| return stats |
|
|
