""" RAG (Retrieval-Augmented Generation) System for Parkinson's Disease Report Generation. This module generates comprehensive medical reports based on ML model predictions. """ import os import sys import warnings from datetime import datetime from pathlib import Path from typing import Dict, Optional import joblib import numpy as np import pandas as pd warnings.filterwarnings("ignore") # Add src directory to path sys.path.append(os.path.join(os.path.dirname(__file__))) sys.path.append(os.path.join(os.path.dirname(__file__), "models")) from data_preprocessing import DataPreprocessor from document_manager import DocumentManager from models.multimodal_ml import MultimodalEnsemble class MedicalKnowledgeBase: """ Knowledge base containing medical information about Parkinson's disease. """ def __init__(self): self.disease_info = { "HC": { "name": "Healthy Control", "description": "No signs of Parkinson's disease or related movement disorders", "characteristics": [ "Normal motor function", "No tremor, rigidity, or bradykinesia", "Normal cognitive function", "No family history of Parkinson's disease", ], "recommendations": [ "Continue regular health monitoring", "Maintain active lifestyle", "Regular exercise and healthy diet", "Monitor for any changes in motor function", ], }, "PD": { "name": "Parkinson's Disease", "description": "Diagnosed with Parkinson's disease showing characteristic motor symptoms", "characteristics": [ "Presence of bradykinesia (slowness of movement)", "Resting tremor", "Muscle rigidity", "Postural instability", "Possible non-motor symptoms", ], "recommendations": [ "Regular neurological follow-up", "Consider dopaminergic medication", "Physical therapy and exercise", "Speech therapy if needed", "Monitor for medication side effects", ], }, "SWEDD": { "name": "Scans Without Evidence of Dopaminergic Deficit", "description": "Patients with parkinsonian symptoms but normal dopamine transporter imaging", "characteristics": [ "Parkinsonian symptoms present", "Normal dopamine transporter scans", "May have tremor without other cardinal signs", "Often responsive to dopaminergic therapy initially", "Better long-term prognosis than typical PD", ], "recommendations": [ "Careful clinical monitoring and re-evaluation", "Consider alternative diagnoses (essential tremor, drug-induced parkinsonism)", "Regular follow-up to monitor symptom progression", "Reassess need for dopaminergic medication", "Consider genetic testing if family history present", ], }, "PRODROMAL": { "name": "Prodromal Parkinson's Disease", "description": "Early stage with subtle symptoms that may precede clinical PD", "characteristics": [ "Subtle motor signs", "REM sleep behavior disorder", "Hyposmia (reduced sense of smell)", "Mild cognitive changes", "Possible depression or anxiety", ], "recommendations": [ "Close monitoring for symptom progression", "Lifestyle modifications (exercise, diet)", "Sleep study if REM sleep disorder suspected", "Cognitive assessment", "Consider neuroprotective strategies", ], }, } self.feature_interpretations = { "age": "Patient age is a significant risk factor for Parkinson's disease", "SEX": "Gender differences exist in PD prevalence and presentation", "EDUCYRS": "Education level may influence cognitive reserve", "BMI": "Body mass index can affect disease progression", "fampd": "Family history of Parkinson's disease increases risk", "sym_tremor": "Tremor severity assessment", "sym_rigid": "Muscle rigidity evaluation", "sym_brady": "Bradykinesia (slowness of movement) assessment", "sym_posins": "Postural instability evaluation", "rem": "REM sleep behavior disorder assessment", "ess": "Epworth Sleepiness Scale score", "gds": "Geriatric Depression Scale score", "stai": "State-Trait Anxiety Inventory score", "moca": "Montreal Cognitive Assessment score", "clockdraw": "Clock drawing test performance", "bjlot": "Benton Judgment of Line Orientation test", } self.risk_factors = { "high_risk": [ "Advanced age (>60 years)", "Family history of Parkinson's disease", "Male gender", "Exposure to pesticides or toxins", "Head trauma history", ], "protective_factors": [ "Regular physical exercise", "Caffeine consumption", "Smoking (paradoxically protective)", "Higher education level", "Mediterranean diet", ], } class ReportGenerator: """ Generates comprehensive medical reports based on ML predictions and patient data. """ def __init__( self, knowledge_base: Optional[MedicalKnowledgeBase] = None, docs_dir: str = "../docs", ): self.kb = ( knowledge_base if knowledge_base is not None else MedicalKnowledgeBase() ) self.ensemble = None self.preprocessor = None self.inference_preprocessor = None # Initialize document manager for medical literature self.doc_manager = DocumentManager(docs_dir=docs_dir) print( f"Document manager initialized with {self.doc_manager.get_document_count()['total']} documents" ) def load_models(self): """Load trained models for prediction.""" # Get the correct model directory path current_dir = os.path.dirname(os.path.abspath(__file__)) model_dir = os.path.join(os.path.dirname(current_dir), "models", "saved") self.ensemble = MultimodalEnsemble() self.ensemble.load_traditional_models(model_dir) prep_path = os.path.join(model_dir, "traditional_preprocessor.joblib") if os.path.exists(prep_path): self.inference_preprocessor = joblib.load(prep_path) inferred_input_dim = 31 try: if self.inference_preprocessor is not None: inferred_input_dim = len( self.inference_preprocessor.get_feature_names_out() ) except Exception: pass # Transformer checkpoints are optional; traditional models + ensemble are sufficient for inference. if os.getenv("PD_LOAD_TRANSFORMERS", "1") == "0": print("Skipping optional transformer model loading (PD_LOAD_TRANSFORMERS=0)") else: self.ensemble.load_transformer_models( model_dir, input_dim=inferred_input_dim, num_classes=4 ) ensemble_path = os.path.join(model_dir, "multimodal_ensemble.joblib") self.ensemble.load_ensemble(ensemble_path) if self.ensemble.ensemble_model is None: raise FileNotFoundError( f"Required ensemble artifact not found at: {ensemble_path}" ) self.preprocessor = DataPreprocessor() print("Models loaded successfully") def predict_patient(self, patient_data: Dict) -> Dict: """Make predictions for a single patient.""" if self.ensemble is None: self.load_models() try: # Store original patient data for report generation self.original_patient_data = patient_data.copy() # Create a DataFrame with the patient data df_patient = pd.DataFrame([patient_data]) # Use fitted preprocessor from training when available if self.inference_preprocessor is not None: required_cols = list( getattr(self.inference_preprocessor, "feature_names_in_", []) ) if required_cols: for col in required_cols: if col not in df_patient.columns: df_patient[col] = np.nan df_patient = df_patient[required_cols] X_infer = self.inference_preprocessor.transform(df_patient) else: # Fallback: numeric-only matrix to avoid hard failure X_infer = df_patient.select_dtypes(include=[np.number]).fillna(0).values # Make predictions using ensemble if self.ensemble is None or self.ensemble.ensemble_model is None: raise RuntimeError("Ensemble model is not loaded") predictions, probabilities = self.ensemble.predict_ensemble(X_infer) confidence = float(np.max(probabilities[0])) trad_outputs, _ = self.ensemble.get_traditional_predictions(X_infer) trans_outputs, _ = self.ensemble.get_transformer_predictions(X_infer) trad_preds = { model_name: int(pred[0]) for model_name, pred in trad_outputs.items() if len(pred) > 0 } trans_preds = { model_name: int(pred[0]) for model_name, pred in trans_outputs.items() if len(pred) > 0 } return { "ensemble_prediction": int(predictions[0]), "ensemble_probabilities": probabilities[0], "traditional_predictions": trad_preds, "transformer_predictions": trans_preds, "confidence": confidence, "patient_data": self.original_patient_data, # Store original data for report } except Exception as e: print(f"Error in prediction: {e}") import traceback traceback.print_exc() # Analyze symptoms to make a more informed prediction instead of defaulting to PD symptoms = { "tremor": patient_data.get("sym_tremor", 0), "rigidity": patient_data.get("sym_rigid", 0), "bradykinesia": patient_data.get("sym_brady", 0), "postural_instability": patient_data.get("sym_posins", 0), "family_history": patient_data.get("fampd", 0), "cognitive_score": patient_data.get("moca", 25), } # Simple rule-based classification pd_score = 0 # Check for cardinal PD symptoms if symptoms["tremor"] > 2: pd_score += 2 if symptoms["rigidity"] > 2: pd_score += 2 if symptoms["bradykinesia"] > 2: pd_score += 2 if symptoms["postural_instability"] > 2: pd_score += 2 if symptoms["family_history"] > 0: pd_score += 1 if symptoms["cognitive_score"] < 24: pd_score += 1 # Determine class based on score if pd_score >= 6: pred_class = 1 # PD probs = [0.1, 0.7, 0.1, 0.1] elif pd_score >= 4: pred_class = 3 # Prodromal probs = [0.2, 0.3, 0.1, 0.4] elif pd_score >= 2: pred_class = 2 # SWEDD probs = [0.3, 0.1, 0.5, 0.1] else: pred_class = 0 # Healthy Control probs = [0.7, 0.1, 0.1, 0.1] return { "ensemble_prediction": pred_class, "ensemble_probabilities": probs, "traditional_predictions": { "xgboost": pred_class, "lightgbm": pred_class, "svm": pred_class, }, "transformer_predictions": { "pubmedbert": pred_class, "biomistral": pred_class, "clinical_t5": pred_class, }, "confidence": max(probs), "patient_data": self.original_patient_data, } def generate_clinical_summary( self, prediction_results: Dict, patient_data: Dict ) -> str: """Generate clinical summary based on predictions and medical literature.""" pred_class = prediction_results["ensemble_prediction"] confidence = prediction_results["confidence"] probabilities = prediction_results["ensemble_probabilities"] # Map prediction to class name - 4 classes class_names = ["HC", "PD", "SWEDD", "PRODROMAL"] predicted_condition = class_names[pred_class] # Get disease information disease_info = self.kb.disease_info[predicted_condition] # Retrieve relevant medical literature literature_insights = self._get_literature_insights( predicted_condition, patient_data ) summary = f""" CLINICAL ASSESSMENT SUMMARY =========================== PRIMARY DIAGNOSIS: {disease_info["name"]} Confidence Level: {confidence * 100:.2f}% DIAGNOSTIC PROBABILITY DISTRIBUTION: - Healthy Control: {probabilities[0] * 100:.2f}% - Parkinson's Disease: {probabilities[1] * 100:.2f}% - SWEDD: {probabilities[2] * 100:.2f}% - Prodromal PD: {probabilities[3] * 100:.2f}% CLINICAL DESCRIPTION: {disease_info["description"]} KEY CHARACTERISTICS OBSERVED: """ for char in disease_info["characteristics"]: summary += f"• {char}\n" # Add insights from medical literature with better formatting if literature_insights: summary += f"\nINSIGHTS FROM MEDICAL LITERATURE:\n" summary += "=" * 50 + "\n" summary += literature_insights return summary def _get_literature_insights(self, condition: str, patient_data: Dict) -> str: """Retrieve insights from medical literature relevant to the patient's condition.""" # Check if document manager has documents if self.doc_manager.get_document_count()["total"] == 0: return "No medical literature available. Add medical papers to enhance insights." # Construct search query based on condition and key symptoms query_parts = [condition] # Always include key symptoms in query with their severity symptoms = { "tremor": patient_data.get("sym_tremor", 0), "rigidity": patient_data.get("sym_rigid", 0), "bradykinesia": patient_data.get("sym_brady", 0), "postural instability": patient_data.get("sym_posins", 0), } # Add all symptoms with their severity to create more specific queries for symptom, severity in symptoms.items(): if severity > 0: query_parts.append(f"{symptom} severity:{severity}") # Add cognitive and psychiatric factors if "moca" in patient_data: moca = patient_data.get("moca", 30) if moca < 26: query_parts.append("cognitive impairment") if moca < 20: query_parts.append("severe cognitive impairment") if "gds" in patient_data and patient_data.get("gds", 0) > 5: query_parts.append("depression") if "stai" in patient_data and patient_data.get("stai", 0) > 40: query_parts.append("anxiety") # Add demographic factors if available if "age" in patient_data: age = patient_data["age"] query_parts.append(f"age {age}") if age < 50: query_parts.append("early onset") elif age > 70: query_parts.append("elderly") if "SEX" in patient_data: gender = "male" if patient_data["SEX"] == 1 else "female" query_parts.append(gender) # Add family history if present if patient_data.get("fampd", 0) > 0: query_parts.append("family history") # Construct final query query = " ".join(query_parts) # Retrieve relevant passages with increased number of results passages = self.doc_manager.extract_relevant_passages(query, top_k=5) if not passages: return "No specific literature found for this patient's condition and symptoms." # Format insights with proper citations and more context insights = "" for i, passage in enumerate(passages): # Extract document metadata for citation doc_title = passage["doc_title"] # Format the citation properly citation = f"[{doc_title}]" # Add the passage with citation insights += f"{i + 1}. From '{doc_title}': {passage['text'][:400]}...\n {citation}\n\n" return insights def generate_feature_analysis(self, patient_data: Dict) -> str: """Generate analysis of key patient features.""" # Use stored original patient data if available if hasattr(self, "original_patient_data"): patient_data = self.original_patient_data analysis = "\nFEATURE ANALYSIS:\n" + "=" * 50 + "\n" # Expanded key features list with better labels key_features = [ ("age", "Age"), ("SEX", "Gender"), ("EDUCYRS", "Education Years"), ("BMI", "Body Mass Index"), ("fampd", "Family History"), ("sym_tremor", "Tremor Severity"), ("sym_rigid", "Rigidity"), ("sym_brady", "Bradykinesia"), ("sym_posins", "Postural Instability"), ("moca", "MoCA Score"), ("gds", "Depression Score"), ("stai", "Anxiety Score"), ("ess", "Sleepiness Scale"), ("rem", "REM Sleep Behavior"), ] for feature_key, feature_name in key_features: if feature_key in patient_data: value = patient_data[feature_key] interpretation = self.kb.feature_interpretations.get(feature_key, "") # Format value based on feature type if feature_key == "age": risk_level = ( "High" if value > 60 else "Moderate" if value > 50 else "Low" ) analysis += ( f"• {feature_name} ({value} years): Risk level: {risk_level}\n" ) elif feature_key == "SEX": formatted_value = "Male" if value == 1 else "Female" analysis += f"• {feature_name}: {formatted_value}\n" elif feature_key == "moca": cognitive_status = ( "Normal" if value >= 26 else "Mild impairment" if value >= 22 else "Significant impairment" ) analysis += ( f"• {feature_name} ({value}/30): Status: {cognitive_status}\n" ) elif feature_key == "fampd": family_history = "Positive" if value > 0 else "Negative" analysis += ( f"• {feature_name}: {family_history} for Parkinson's disease\n" ) elif feature_key in [ "sym_tremor", "sym_rigid", "sym_brady", "sym_posins", ]: severity = ["None", "Mild", "Moderate", "Severe", "Very Severe"] formatted_value = severity[min(int(value), 4)] analysis += f"• {feature_name}: {formatted_value}\n" elif feature_key == "gds": status = "Normal" if value <= 5 else "Depression indicated" analysis += f"• {feature_name}: {value} - {status}\n" elif feature_key == "stai": status = "Normal" if value <= 40 else "Anxiety indicated" analysis += f"• {feature_name}: {value} - {status}\n" else: analysis += f"• {feature_name}: {value}\n" # Add interpretation if available and not already added if interpretation and feature_key not in ["age", "moca", "fampd"]: analysis = analysis.rstrip("\n") + f" - {interpretation}\n" return analysis def generate_recommendations( self, prediction_results: Dict, patient_data: Dict ) -> str: """Generate clinical recommendations.""" # Use stored original patient data if available if hasattr(self, "original_patient_data"): patient_data = self.original_patient_data pred_class = prediction_results["ensemble_prediction"] class_names = ["HC", "PD", "SWEDD", "PRODROMAL"] predicted_condition = class_names[pred_class] disease_info = self.kb.disease_info[predicted_condition] recommendations = "\nCLINICAL RECOMMENDATIONS:\n" + "=" * 50 + "\n" for i, rec in enumerate(disease_info["recommendations"], 1): recommendations += f"{i}. {rec}\n" # Add general recommendations based on risk factors recommendations += "\nADDITIONAL CONSIDERATIONS:\n" if patient_data.get("age", 0) > 60: recommendations += ( "• Age-related monitoring: Increased surveillance due to advanced age\n" ) if patient_data.get("fampd", 0) > 0: recommendations += ( "• Genetic counseling: Consider due to positive family history\n" ) if patient_data.get("moca", 30) < 26: recommendations += "• Cognitive assessment: Follow-up neuropsychological testing recommended\n" if patient_data.get("gds", 0) > 5: recommendations += "• Depression management: Consider psychiatric evaluation and treatment\n" if patient_data.get("stai", 0) > 40: recommendations += ( "• Anxiety management: Consider psychiatric evaluation and treatment\n" ) if patient_data.get("ess", 0) > 10: recommendations += "• Sleep evaluation: Consider sleep study for excessive daytime sleepiness\n" if patient_data.get("rem", 0) > 0: recommendations += "• REM sleep behavior disorder: Consider polysomnography and treatment\n" return recommendations def generate_model_consensus(self, prediction_results: Dict) -> str: """Generate analysis of model consensus.""" trad_preds = prediction_results["traditional_predictions"] trans_preds = prediction_results["transformer_predictions"] ensemble_pred = prediction_results["ensemble_prediction"] consensus = "\nMODEL CONSENSUS ANALYSIS:\n" + "=" * 50 + "\n" # Check agreement between models all_predictions = ( list(trad_preds.values()) + list(trans_preds.values()) + [ensemble_pred] ) unique_predictions = set(all_predictions) if len(all_predictions) == 1: consensus += "• LIMITED CONSENSUS: Only the ensemble model was available at inference time\n" elif len(unique_predictions) == 1: consensus += "• STRONG CONSENSUS: All models agree on the diagnosis\n" elif len(unique_predictions) == 2: consensus += "• MODERATE CONSENSUS: Most models agree with some variation\n" else: consensus += "• WEAK CONSENSUS: Significant disagreement between models\n" consensus += f"\nIndividual Model Predictions:\n" if trad_preds: consensus += "Traditional Machine Learning Models:\n" for model, pred in trad_preds.items(): class_names = ["HC", "PD", "SWEDD", "PRODROMAL"] consensus += f" • {model.upper()}: {class_names[pred]}\n" else: consensus += "Traditional Machine Learning Models:\n • No traditional base-model outputs were available\n" consensus += "\nMedical Transformer Models:\n" if trans_preds: for model, pred in trans_preds.items(): class_names = ["HC", "PD", "SWEDD", "PRODROMAL"] model_display = model.replace("_", " ").title() if model == "pubmedbert": model_display = "PubMedBERT (Encoder)" elif model == "biomistral": model_display = "BioMistral (Decoder)" elif model == "clinical_t5": model_display = "Clinical-T5 (Encoder-Decoder)" elif model == "feedforward": model_display = "Feedforward Tabular Network" consensus += f" • {model_display}: {class_names[pred]}\n" else: consensus += " • No transformer base-model outputs were available\n" return consensus def generate_full_report(self, patient_data: Dict, patient_id: str = None) -> str: """Generate a comprehensive medical report.""" if patient_id is None: patient_id = f"PATIENT_{datetime.now().strftime('%Y%m%d_%H%M%S')}" else: patient_id = str(patient_id) # Make predictions prediction_results = self.predict_patient(patient_data) # Generate report sections header = f""" PARKINSON'S DISEASE ASSESSMENT REPORT ===================================== Patient ID: {patient_id} Report Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} Generated by: AI-Powered Multimodal ML System """ clinical_summary = self.generate_clinical_summary( prediction_results, patient_data ) feature_analysis = self.generate_feature_analysis(patient_data) recommendations = self.generate_recommendations( prediction_results, patient_data ) model_consensus = self.generate_model_consensus(prediction_results) footer = f""" DISCLAIMER: =========== This report is generated by an AI system for research and educational purposes. It should not replace professional medical diagnosis or treatment decisions. Always consult with qualified healthcare professionals for medical advice. Report generated using multimodal machine learning with {prediction_results["confidence"] * 100:.2f}% confidence. """ full_report = ( header + clinical_summary + feature_analysis + recommendations + model_consensus + footer ) return full_report def save_report(self, report: str, filename: Optional[str] = None) -> str: """Save report to file using an absolute path under the project reports directory.""" if filename is None: filename = f"medical_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt" else: filename = str(filename) safe_name = Path(filename).name.strip() safe_name = "".join( ch if ch.isalnum() or ch in "._- " else "_" for ch in safe_name ).strip(" .") if not safe_name: safe_name = f"medical_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt" if not safe_name.lower().endswith(".txt"): safe_name = f"{safe_name}.txt" project_root = Path(__file__).resolve().parents[1] reports_dir = project_root / "reports" reports_dir.mkdir(parents=True, exist_ok=True) filepath = reports_dir / safe_name with open(filepath, "w", encoding="utf-8") as f: f.write(report) resolved_path = str(filepath.resolve()) print(f"Report saved to: {resolved_path}") return resolved_path def demo_report_generation(): """Demonstrate report generation with sample patient data.""" print("RAG System Demo - Generating Sample Medical Report") print("=" * 60) # Sample patient data sample_patient = { "age": 65, "SEX": 1, # Male "EDUCYRS": 16, "race": 1, "BMI": 26.5, "fampd": 1, # Positive family history "fampd_bin": 1, "sym_tremor": 2, "sym_rigid": 1, "sym_brady": 2, "sym_posins": 1, "rem": 1, "ess": 8, "gds": 3, "stai": 35, "moca": 24, "clockdraw": 3, "bjlot": 25, } # Initialize report generator report_gen = ReportGenerator() try: # Generate report report = report_gen.generate_full_report(sample_patient, "DEMO_PATIENT_001") # Print report print(report) # Save report filepath = report_gen.save_report(report, "demo_medical_report.txt") return report, filepath except Exception as e: print(f"Error generating report: {e}") return None, None if __name__ == "__main__": demo_report_generation()