Upload 2 files

app_test.py

@@ -0,0 +1,371 @@
+import streamlit as st
+import torch
+import numpy as np
+import pickle
+import json
+from transformers import AutoTokenizer, AutoModel
+import torch.nn as nn
+import os
+
+# Set page config
+st.set_page_config(
+    page_title="Drug Prediction and Polypharmacy System", 
+    page_icon="💊",
+    layout="wide"
+)
+
+# Model class definition - must match the training model architecture
+class EnhancedMedicationModel(nn.Module):
+    def __init__(self, model_name, num_medications, num_polypharmacy_classes, num_disease_classes, dropout_rate=0.3):
+        super().__init__()
+        self.bert = AutoModel.from_pretrained(model_name)
+        self.dropout = nn.Dropout(dropout_rate)
+        hidden_size = self.bert.config.hidden_size
+
+        # Common representation layer
+        self.common_dense = nn.Linear(hidden_size, hidden_size)
+
+        # Task-specific layers with increased complexity
+        # Medication prediction head (multi-label)
+        self.medication_classifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size//2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size//2, num_medications)
+        )
+
+        # Polypharmacy risk head (multi-class)
+        self.polypharmacy_classifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size//2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size//2, num_polypharmacy_classes)
+        )
+
+        # Disease prediction head (multi-class)
+        self.disease_classifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size//2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size//2, num_disease_classes)
+        )
+
+        # Apply weight initialization
+        self._init_weights()
+
+    def _init_weights(self):
+        # Initialize weights for better convergence
+        for module in [self.medication_classifier, self.polypharmacy_classifier,
+                      self.disease_classifier, self.common_dense]:
+            if isinstance(module, nn.Sequential):
+                for layer in module:
+                    if isinstance(layer, nn.Linear):
+                        nn.init.xavier_normal_(layer.weight)
+                        nn.init.zeros_(layer.bias)
+            elif isinstance(module, nn.Linear):
+                nn.init.xavier_normal_(module.weight)
+                nn.init.zeros_(layer.bias)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.last_hidden_state[:, 0, :]  # CLS token
+        pooled_output = self.dropout(pooled_output)
+
+        # Common representation
+        common_features = torch.relu(self.common_dense(pooled_output))
+
+        medication_logits = self.medication_classifier(common_features)
+        polypharmacy_logits = self.polypharmacy_classifier(common_features)
+        disease_logits = self.disease_classifier(common_features)
+
+        return medication_logits, polypharmacy_logits, disease_logits
+
+@st.cache_resource
+def load_model_and_resources():
+    """Load model and necessary resources (cached for performance)"""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Load model configuration - fixed file paths
+    with open('streamlit_model/model_config.json', 'r') as f:
+        model_config = json.load(f)
+    
+    # Initialize model
+    model_name = model_config['model_name']
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    
+    # Create model architecture
+    model = EnhancedMedicationModel(
+        model_name=model_name,
+        num_medications=model_config['num_medications'],
+        num_polypharmacy_classes=model_config['num_polypharmacy_classes'],
+        num_disease_classes=model_config['num_disease_classes'],
+        dropout_rate=0.3
+    )
+    
+    # Load trained weights - fixed file path
+    model.load_state_dict(torch.load('streamlit_model/model_state_dict.pt', map_location=device))
+    model = model.to(device)
+    model.eval()
+    
+    # Load encoders - fixed file path
+    with open('streamlit_model/label_encoders.pkl', 'rb') as f:
+        encoders = pickle.load(f)
+    
+    # Load lookup data - fixed file path
+    with open('streamlit_model/lookup_data.pkl', 'rb') as f:
+        lookup_data = pickle.load(f)
+    
+    return {
+        'model': model,
+        'tokenizer': tokenizer,
+        'mlb': encoders['mlb'],
+        'le_risk': encoders['le_risk'],
+        'le_disease': encoders['le_disease'],
+        'lookup_data': lookup_data,
+        'device': device
+    }
+
+def predict_patient_health_profile(patient_data, resources):
+    """
+    Predict health profile for a patient based on input data
+    """
+    model = resources['model']
+    tokenizer = resources['tokenizer']
+    mlb = resources['mlb']
+    le_risk = resources['le_risk']
+    le_disease = resources['le_disease']
+    lookup_data = resources['lookup_data']
+    device = resources['device']
+    
+    # Create text input
+    text_input = f"Patient age {patient_data['age']}, gender {patient_data['gender']}, blood group {patient_data['blood_group']}, weight {patient_data['weight']}kg. " +                 f"SYMPTOMS: {patient_data['symptoms']}. " +                 f"SEVERITY: {patient_data['severity']}."
+    
+    # Tokenize
+    encoding = tokenizer(
+        text_input,
+        add_special_tokens=True,
+        max_length=256,
+        padding='max_length',
+        truncation=True,
+        return_tensors='pt'
+    )
+    
+    # Move to device
+    input_ids = encoding['input_ids'].to(device)
+    attention_mask = encoding['attention_mask'].to(device)
+    
+    # Get predictions
+    with torch.no_grad():
+        medication_logits, polypharmacy_logits, disease_logits = model(input_ids, attention_mask)
+        medication_preds = torch.sigmoid(medication_logits) > 0.5
+        polypharmacy_pred = torch.argmax(polypharmacy_logits, dim=1)
+        disease_pred = torch.argmax(disease_logits, dim=1)
+    
+    # Convert predictions to human-readable format
+    predicted_medications = mlb.classes_[medication_preds[0].cpu().numpy()]
+    predicted_risk = le_risk.classes_[polypharmacy_pred.item()]
+    predicted_disease = le_disease.classes_[disease_pred.item()]
+    
+    # Get medication probabilities for all medications
+    medication_probs = torch.sigmoid(medication_logits).cpu().numpy()[0]
+    med_prob_dict = {med: prob for med, prob in zip(mlb.classes_, medication_probs)}
+    
+    # Sort medications by probability
+    sorted_meds = sorted(med_prob_dict.items(), key=lambda x: x[1], reverse=True)
+    top_meds = sorted_meds[:5]  # Get top 5 medications
+    
+    # Format medication results
+    med_results = []
+    for i, med in enumerate(predicted_medications[:3]):
+        med_details = {
+            'medication': med,
+            'dosage': 'Consult doctor',
+            'frequency': 'Consult doctor',
+            'instruction': 'Consult doctor',
+            'duration': 'As prescribed',
+            'confidence': float(med_prob_dict[med])
+        }
+        med_results.append(med_details)
+    
+    # Get disease information
+    disease_causes = lookup_data['disease_causes_dict'].get(predicted_disease, "Unknown causes")
+    disease_prevention = lookup_data['disease_prevention_dict'].get(predicted_disease, "Consult healthcare provider")
+    
+    # Get polypharmacy recommendation
+    polypharmacy_recommendation = lookup_data['polypharmacy_recommendation_dict'].get(
+        predicted_risk, "Consult healthcare provider"
+    )
+    
+    # Get personalized health tip
+    age_decade = (patient_data['age'] // 10) * 10
+    health_tip_key = (predicted_disease, age_decade, patient_data['gender'])
+    personalized_health_tip = lookup_data['health_tips_dict'].get(
+        health_tip_key, "Maintain a balanced diet and regular exercise routine."
+    )
+    
+    # Return comprehensive results
+    return {
+        'patient_name': patient_data['name'],  # Include patient name in results
+        'predicted_disease': predicted_disease,
+        'disease_causes': disease_causes,
+        'disease_prevention': disease_prevention,
+        'medications': med_results,
+        'polypharmacy_risk': predicted_risk,
+        'polypharmacy_recommendation': polypharmacy_recommendation,
+        'personalized_health_tips': personalized_health_tip,
+        'medication_probabilities': {med: float(prob) for med, prob in top_meds}
+    }
+
+def main():
+    # App title and description
+    st.title("🏥 Drug Prediction and Polypharmacy System")
+    st.markdown("Enter patient information to receive medication recommendations, disease prediction, and polypharmacy risk assessment.")
+    
+    try:
+        # Load model and resources
+        with st.spinner("Loading medical model and resources..."):
+            resources = load_model_and_resources()
+        
+        # Create two columns for input form
+        col1, col2 = st.columns(2)
+        
+        # Patient information inputs
+        with col1:
+            st.subheader("Patient Information")
+            # Add patient name input field
+            name = st.text_input("Patient Name", value="John Doe")
+            age = st.number_input("Age", min_value=1, max_value=120, value=45)
+            gender = st.selectbox("Gender", options=["Male", "Female", "Other"])
+            blood_group = st.selectbox("Blood Group", options=["A+", "A-", "B+", "B-", "AB+", "AB-", "O+", "O-"])
+            weight = st.number_input("Weight (kg)", min_value=1.0, max_value=300.0, value=70.0, step=0.1)
+        
+        with col2:
+            st.subheader("Symptoms Information")
+            
+            # Common symptoms options
+            common_symptoms = [
+                "Headache", "Fever", "Fatigue", "Nausea", "Cough", 
+                "Sore throat", "Shortness of breath", "Chest pain", 
+                "Dizziness", "Abdominal pain", "Vomiting", "Diarrhea",
+                "Muscle ache", "Joint pain", "Rash", "Loss of appetite"
+            ]
+            
+            # Use multiselect for symptoms selection
+            selected_symptoms = st.multiselect(
+                "Select Symptoms", 
+                options=common_symptoms,
+                default=["Headache", "Fever", "Fatigue"]
+            )
+            
+            # Custom symptom input
+            custom_symptom = st.text_input("Add other symptom (if not in list)")
+            if custom_symptom:
+                selected_symptoms.append(custom_symptom)
+            
+            # Convert selected symptoms to string format as expected by the model
+            symptoms = "; ".join(selected_symptoms)
+            
+            # More compact severity selection
+            st.subheader("Symptom Severity")
+            
+            # Define severity levels
+            severity_levels = {
+                "Very Mild": 1,
+                "Mild": 2,
+                "Moderate": 3,
+                "Severe": 4,
+                "Very Severe": 5
+            }
+            
+            severity_dict = {}
+            
+            # Create a more compact layout with 2 columns for severity selection
+            if selected_symptoms:
+                cols = st.columns(2)
+                for i, symptom in enumerate(selected_symptoms):
+                    # Alternate between columns
+                    with cols[i % 2]:
+                        severity_option = st.selectbox(
+                            f"{symptom}",
+                            options=list(severity_levels.keys()),
+                            index=1  # Default to "Mild"
+                        )
+                        severity_dict[symptom] = severity_levels[severity_option]
+            
+            # Convert severity dict to string format as expected by the model
+            severity = "; ".join([f"{symptom}:{score}" for symptom, score in severity_dict.items()])
+            
+        # Submit button
+        if st.button("Generate Health Profile", type="primary"):
+            with st.spinner("Analyzing patient data and generating health profile..."):
+                # Prepare patient data
+                patient_data = {
+                    'name': name,  # Include name in patient data
+                    'age': age,
+                    'gender': gender,
+                    'blood_group': blood_group,
+                    'weight': weight,
+                    'symptoms': symptoms,
+                    'severity': severity
+                }
+                
+                # Get prediction
+                prediction = predict_patient_health_profile(patient_data, resources)
+                
+                # Display results in three columns
+                st.subheader(f"🔍 Health Profile Analysis Results for {prediction['patient_name']}")
+                
+                col1, col2, col3 = st.columns([1, 1, 1])
+                
+                # Column 1: Disease information
+                with col1:
+                    st.markdown("### 🦠 Disease Prediction")
+                    st.markdown(f"**Predicted Disease**: {prediction['predicted_disease']}")
+                    
+                    with st.expander("Disease Causes"):
+                        st.write(prediction['disease_causes'])
+                    
+                    with st.expander("Prevention Methods"):
+                        st.write(prediction['disease_prevention'])
+                
+                # Column 2: Medication recommendations
+                with col2:
+                    st.markdown("### 💊 Medication Recommendations")
+                    for i, med in enumerate(prediction['medications']):
+                        st.markdown(f"**{i+1}. {med['medication']}** (Confidence: {med['confidence']:.2f})")
+                        med_details = f"""
+                        - **Dosage:** {med['dosage']}
+                        - **Frequency:** {med['frequency']}
+                        - **Instructions:** {med['instruction']}
+                        - **Duration:** {med['duration']}
+                        """
+                        st.markdown(med_details)
+                        st.divider()
+                
+                # Column 3: Risk assessment and health tips
+                with col3:
+                    st.markdown("### ⚠️ Polypharmacy Assessment")
+                    risk_color = "green" if prediction['polypharmacy_risk'] == "Low" else                                 "orange" if prediction['polypharmacy_risk'] == "Medium" else "red"
+                    st.markdown(f"**Risk Level**: <span style='color:{risk_color};font-weight:bold;'>{prediction['polypharmacy_risk']}</span>", 
+                                unsafe_allow_html=True)
+                    st.markdown(f"**Recommendation**: {prediction['polypharmacy_recommendation']}")
+                    
+                    st.markdown("### 🌿 Personalized Health Tips")
+                    st.info(prediction['personalized_health_tips'])
+                    
+                # Display medication probabilities as text with progress bars
+                st.subheader("Medication Confidence Scores")
+                med_names = list(prediction['medication_probabilities'].keys())
+                med_probs = list(prediction['medication_probabilities'].values())
+                
+                # Display each medication with its confidence score as text and progress bar
+                for med_name, med_prob in zip(med_names, med_probs):
+                    st.text(f"{med_name}: {med_prob:.2f}")
+                    st.progress(med_prob)
+                    
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")
+        st.error("Please make sure all model files are correctly placed in the 'streamlit_model' directory")
+
+if __name__ == "__main__":
+    main()