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"]) # 4180186 = English 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", "") # Filter target vocabularies and only valid/active concepts 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", ""), } # Add canonical concept name as initial synonym concept_synonyms[cid].add(cname) # For SNOMED, also add tag-stripped version if applicable 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 SNOMED, also add tag-stripped synonym 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:,}") # Filter concepts that meet minimum synonym requirement for metric learning 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 requested, sample randomly 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.") # Sort and shuffle deterministically trainable_concept_ids.sort() random.shuffle(trainable_concept_ids) # Concept-level split to prevent data leakage 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):,}") # Write datasets 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: # Prioritize canonical name, sample remaining 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)") # Save metadata lookup dictionary for inference & evaluation 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}") # Save summary statistics 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