DCE / decision /engine /training_data_generator.py
That guy James Bond :)
Deploy Medical Intent Escalation API
af61b34
Raw
History Blame Contribute Delete
13 kB
"""
Training Data Generator (Phase 3C).
Generates synthetic training samples from taxonomy trigger definitions
for model retraining and augmentation.
Each domain's triggers.yaml contains carefully curated phrases at
multiple confidence tiers. These can be converted into labeled training
examples to improve DriveHealthBERT's domain coverage.
Capabilities:
- Generate ESCALATION samples from R2/R3 domain triggers
- Generate domain-specific samples with label mapping
- Generate negation examples (negative samples for safety)
- Template-based augmentation with variations
- JSONL output compatible with train.py
Usage:
generator = TrainingDataGenerator(config)
samples = generator.generate_all()
generator.write_jsonl(samples, "data/taxonomy_augmented.jsonl")
"""
from __future__ import annotations
import json
import logging
import random
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple
from decision.engine.config_loader import DecisionConfigLoader
logger = logging.getLogger("decision.training_data_generator")
# Map domains to DriveHealthBERT labels
_DOMAIN_TO_LABEL = {
# R3 domains → ESCALATION
"chest_pain": "ESCALATION",
"suicidal_ideation": "ESCALATION",
"homicidal_ideation": "ESCALATION",
"stroke_symptoms": "ESCALATION",
"seizure": "ESCALATION",
"severe_bleeding": "ESCALATION",
"anaphylaxis": "ESCALATION",
"airway_compromise": "ESCALATION",
"overdose": "ESCALATION",
"pregnancy_emergency": "ESCALATION",
"ectopic_signal": "ESCALATION",
"febrile_neutropenia": "ESCALATION",
"meningitis_signal": "ESCALATION",
"dka_hhs": "ESCALATION",
"severe_hypoglycemia": "ESCALATION",
"syncope": "ESCALATION",
"thunderclap_headache": "ESCALATION",
"sudden_vision_loss": "ESCALATION",
"acute_limb_deficit": "ESCALATION",
# R2 domains → ESCALATION (moderate, but still escalation)
"shortness_of_breath_rest": "ESCALATION",
"palpitations_dizziness": "ESCALATION",
"heart_failure_worsening": "ESCALATION",
"gi_bleed": "ESCALATION",
"severe_abdominal_pain": "ESCALATION",
"wound_infection": "ESCALATION",
"fall_high_risk": "ESCALATION",
"behavioral_crisis": "ESCALATION",
"acute_confusion": "ESCALATION",
"persistent_vomiting": "ESCALATION",
"post_discharge_fever": "ESCALATION",
"dvt_signal": "ESCALATION",
"npo_violation": "ESCALATION",
"anticoagulant_hold_failure": "ESCALATION",
"acute_illness_preop": "ESCALATION",
"worsening_depression_si": "ESCALATION",
"procedure_anxiety_severe": "ESCALATION",
"domestic_safety_concern": "ESCALATION",
# R1/R0 symptom domains → SYMPTOM_CHECK
"exertional_dyspnea": "SYMPTOM_CHECK",
"persistent_dizziness": "SYMPTOM_CHECK",
"abnormal_bp": "SYMPTOM_CHECK",
"reproducible_chest_discomfort": "SYMPTOM_CHECK",
"wound_concern": "SYMPTOM_CHECK",
"uncontrolled_pain": "SYMPTOM_CHECK",
"worsening_chronic_pain": "SYMPTOM_CHECK",
"progressive_weakness": "SYMPTOM_CHECK",
"urinary_retention": "SYMPTOM_CHECK",
"severe_diarrhea": "SYMPTOM_CHECK",
"poor_oral_intake": "SYMPTOM_CHECK",
"new_jaundice": "SYMPTOM_CHECK",
"glucose_out_of_range": "SYMPTOM_CHECK",
"recurrent_falls": "SYMPTOM_CHECK",
"post_procedure_incontinence": "SYMPTOM_CHECK",
# Medication domains → MEDICATION
"medication": "MEDICATION",
"medication_nonadherence": "MEDICATION",
# Scheduling domains → APPOINTMENT
"scheduling": "APPOINTMENT",
"scheduling_barrier": "APPOINTMENT",
"missed_followup": "APPOINTMENT",
"post_discharge_engagement": "APPOINTMENT",
# Screening domains → SYMPTOM_CHECK
"phq2_positive": "SYMPTOM_CHECK",
"audit_c_positive": "SYMPTOM_CHECK",
"screening_gap_identified": "SYMPTOM_CHECK",
"screening_refusal": "GENERAL_INQUIRY",
"cognitive_concern": "SYMPTOM_CHECK",
# Social/safety domains
"caregiver_safety_concern": "ESCALATION",
"food_insecurity": "GENERAL_INQUIRY",
"social_isolation": "GENERAL_INQUIRY",
"tobacco_use_active": "GENERAL_INQUIRY",
}
# Templates for wrapping trigger phrases into natural sentences
_PATIENT_TEMPLATES = [
"{phrase}",
"I have {phrase}",
"I'm experiencing {phrase}",
"I've been having {phrase}",
"I think I have {phrase}",
"My {phrase} is getting worse",
"I need help with {phrase}",
"I'm worried about {phrase}",
"I want to talk about {phrase}",
"Can you help me with {phrase}",
]
_NEGATION_TEMPLATES = [
"I don't have {phrase}",
"No {phrase}",
"I'm not experiencing {phrase}",
"{phrase} has gone away",
"I used to have {phrase} but not anymore",
"My mother had {phrase} not me",
]
class TrainingDataGenerator:
"""Generates training data from taxonomy triggers."""
def __init__(self, config: DecisionConfigLoader):
self._config = config
self._triggers = config.taxonomy_triggers
self._rules = config.taxonomy_rules
def generate_all(
self,
include_negations: bool = True,
max_per_domain: int = 50,
augment_templates: bool = True,
seed: int = 42,
) -> List[Dict[str, str]]:
"""
Generate training samples from all taxonomy domains.
Args:
include_negations: Also generate negative samples from negation patterns
max_per_domain: Maximum samples per domain
augment_templates: Apply template-based augmentation
seed: Random seed for reproducibility
Returns:
List of {"text": ..., "label": ..., "source": ...} dicts
"""
random.seed(seed)
all_samples: List[Dict[str, str]] = []
for domain_name, triggers in self._triggers.items():
label = _DOMAIN_TO_LABEL.get(domain_name)
if not label:
logger.debug("Skipping domain %s (no label mapping)", domain_name)
continue
domain_samples = self._generate_for_domain(
domain_name, triggers, label, augment_templates, max_per_domain
)
all_samples.extend(domain_samples)
# Generate negation samples
if include_negations and label == "ESCALATION":
neg_samples = self._generate_negations(
domain_name, triggers, max_per_domain // 3
)
all_samples.extend(neg_samples)
random.shuffle(all_samples)
# Summary
label_counts = {}
for s in all_samples:
label_counts[s["label"]] = label_counts.get(s["label"], 0) + 1
logger.info(
"Generated %d training samples from %d domains: %s",
len(all_samples),
len(self._triggers),
label_counts,
)
return all_samples
def _generate_for_domain(
self,
domain_name: str,
triggers: Dict[str, Any],
label: str,
augment: bool,
max_samples: int,
) -> List[Dict[str, str]]:
"""Generate training samples for a single domain."""
samples = []
lexical = triggers.get("lexical_signals", {})
for tier_name in ("high", "medium", "low"):
tier = lexical.get(tier_name, {})
# Direct phrase samples
for phrase in tier.get("phrases", []):
samples.append({
"text": phrase,
"label": label,
"source": f"taxonomy/{domain_name}/{tier_name}/phrase",
})
# Template augmentation
if augment:
templates = random.sample(
_PATIENT_TEMPLATES, min(3, len(_PATIENT_TEMPLATES))
)
for template in templates:
try:
augmented = template.format(phrase=phrase.lower())
if augmented != phrase:
samples.append({
"text": augmented,
"label": label,
"source": f"taxonomy/{domain_name}/{tier_name}/augmented",
})
except (KeyError, IndexError):
pass
# Partial samples (lower confidence)
for partial in tier.get("partials", []):
samples.append({
"text": partial,
"label": label,
"source": f"taxonomy/{domain_name}/{tier_name}/partial",
})
# Deduplicate by text
seen = set()
unique = []
for s in samples:
text_lower = s["text"].lower().strip()
if text_lower not in seen:
seen.add(text_lower)
unique.append(s)
# Cap per domain
if len(unique) > max_samples:
unique = random.sample(unique, max_samples)
return unique
def _generate_negations(
self,
domain_name: str,
triggers: Dict[str, Any],
max_samples: int,
) -> List[Dict[str, str]]:
"""
Generate negative samples from negation patterns.
These become NON-ESCALATION training examples to reduce false positives.
"""
samples = []
negation = triggers.get("negation_handling", {})
patterns = negation.get("patterns", [])
# Use negation patterns directly as negative examples
for pattern in patterns:
samples.append({
"text": pattern,
"label": "GENERAL_INQUIRY", # Negated safety = not escalation
"source": f"taxonomy/{domain_name}/negation",
})
# Apply negation templates to domain phrases
lexical = triggers.get("lexical_signals", {})
high_phrases = lexical.get("high", {}).get("phrases", [])
for phrase in high_phrases[:5]: # Top 5 high-confidence phrases
templates = random.sample(
_NEGATION_TEMPLATES, min(2, len(_NEGATION_TEMPLATES))
)
for template in templates:
try:
negated = template.format(phrase=phrase.lower())
samples.append({
"text": negated,
"label": "GENERAL_INQUIRY",
"source": f"taxonomy/{domain_name}/negation_augmented",
})
except (KeyError, IndexError):
pass
if len(samples) > max_samples:
samples = random.sample(samples, max_samples)
return samples
def write_jsonl(
self, samples: List[Dict[str, str]], output_path: str
) -> int:
"""Write samples to JSONL file compatible with train.py."""
path = Path(output_path)
path.parent.mkdir(parents=True, exist_ok=True)
count = 0
with open(path, "w", encoding="utf-8") as f:
for sample in samples:
record = {
"text": sample["text"],
"label": sample["label"],
}
f.write(json.dumps(record, ensure_ascii=False) + "\n")
count += 1
logger.info("Wrote %d samples to %s", count, path)
return count
def get_label_distribution(
self, samples: List[Dict[str, str]]
) -> Dict[str, int]:
"""Get label distribution of generated samples."""
dist: Dict[str, int] = {}
for s in samples:
dist[s["label"]] = dist.get(s["label"], 0) + 1
return dict(sorted(dist.items(), key=lambda x: x[1], reverse=True))
def get_domain_coverage_report(
self, samples: List[Dict[str, str]]
) -> str:
"""Generate a coverage report."""
domain_count: Dict[str, int] = {}
for s in samples:
source = s.get("source", "")
domain = source.split("/")[1] if "/" in source else "unknown"
domain_count[domain] = domain_count.get(domain, 0) + 1
lines = ["=== Training Data Coverage Report ===", ""]
lines.append(f"Total samples: {len(samples)}")
lines.append(f"Domains covered: {len(domain_count)}")
lines.append("")
dist = self.get_label_distribution(samples)
lines.append("Label distribution:")
for label, count in dist.items():
pct = count / len(samples) * 100
lines.append(f" {label:20s}: {count:5d} ({pct:.1f}%)")
lines.append("")
lines.append("Top domains by sample count:")
for domain, count in sorted(domain_count.items(), key=lambda x: x[1], reverse=True)[:15]:
lines.append(f" {domain:35s}: {count:4d}")
return "\n".join(lines)