FM selection and model added.

app.py

@@ -26,14 +26,41 @@ MARITAL_STATUS_CHOICES = list(MARITAL_STATUS_MAP.keys())
 
 MODEL_DIR = "./models"
 
+# def get_available_models():
+#     if not os.path.exists(MODEL_DIR):
+#         os.makedirs(MODEL_DIR) # Create models directory if it doesn't exist
+#         return ["No models found. Please add .joblib models to the 'models' directory."]
+#     models = [f for f in os.listdir(MODEL_DIR) if f.endswith(".joblib")]
+#     if not models:
+#         return ["No models found. Please add .joblib models to the 'models' directory."]
+#     return models
+
 def get_available_models():
     if not os.path.exists(MODEL_DIR):
         os.makedirs(MODEL_DIR) # Create models directory if it doesn't exist
-        return ["No models found. Please add .joblib models to the 'models' directory."]
+        return {"classical": [], "foundation": []}
+    
     models = [f for f in os.listdir(MODEL_DIR) if f.endswith(".joblib")]
     if not models:
-        return ["No models found. Please add .joblib models to the 'models' directory."]
-    return models
+        return {"classical": [], "foundation": []}
+    
+    # Organize models by type and time period
+    model_dict = {
+        "classical": {
+            "diabetes": "Logistic regression_diabetes.joblib",
+            "24mths": "Logistic regression_diabetes_24mths.joblib",
+            "36mths": "Logistic regression_diabetes_36mths.joblib",
+            "48mths": "Logistic regression_diabetes_48mths.joblib"
+        },
+        "foundation": {
+            "diabetes": "FM_Logistic regression_diabetes.joblib",
+            "24mths": "FM_Logistic regression_diabetes_24mths.joblib",
+            "36mths": "FM_Logistic regression_diabetes_36mths.joblib",
+            "48mths": "FM_Logistic regression_diabetes_48mths.joblib"
+        }
+    }
+    
+    return model_dict
 
 # Define all features in the order your model expects them
 # IMPORTANT: This order must match the training data
@@ -60,7 +87,7 @@ EXPECTED_COLUMNS = [
     'hypertriglyceridemia'
 ]
 
-def predict_diabetes(model_name, sex, race, ethnicity, marital_status, Prior_Mean_Glu,
+def predict_diabetes(model_type, time_period, sex, race, ethnicity, marital_status, Prior_Mean_Glu,
        PT_ELX_GRP_1, PT_ELX_GRP_2, PT_ELX_GRP_3, PT_ELX_GRP_4,
        PT_ELX_GRP_5, PT_ELX_GRP_6, PT_ELX_GRP_7, PT_ELX_GRP_8,
        PT_ELX_GRP_9, PT_ELX_GRP_10, PT_ELX_GRP_13, PT_ELX_GRP_14,
@@ -81,8 +108,14 @@ def predict_diabetes(model_name, sex, race, ethnicity, marital_status, Prior_Mea
        oral_contraceptive, cholelithiasis, acute_cholecystitis,
        hypertriglyceridemia):
 
-    if not model_name or "No models found" in model_name:
-        return "Please select a valid model from the 'models/' directory."
+    if not model_type or not time_period:
+        return "Please select both model type and time period."
+
+    model_dict = get_available_models()
+    model_name = model_dict[model_type][time_period]
+    
+    if not model_name:
+        return "Selected model not found. Please check the model type and time period."
 
     model_path = os.path.join(MODEL_DIR, model_name)
     if not os.path.exists(model_path):
@@ -138,11 +171,30 @@ def predict_diabetes(model_name, sex, race, ethnicity, marital_status, Prior_Mea
 
     # Make prediction
     try:
-        prediction = model.predict(df)
-        # You might need to access the first element if prediction is an array
-        # e.g., result = prediction[0] 
-        # Also, convert to a more human-readable output
-        result = prediction[0] 
+        if model_type == "foundation":
+            # Load the TabPFN model for preprocessing
+            try:
+                import numpy as np
+                import tabpfn
+                clf = joblib.load('models/FM/TabPFN_model_chunk_0.joblib')
+                # Get embeddings for the input data
+                X = clf.get_embeddings(df)
+                print(X.shape)
+                # X = np.concatenate(X,axis=1)
+                # X = np.swapaxes(X,0,1)
+                X = X.reshape(768 ,-1)
+                print(X.shape)
+                X = pd.DataFrame(data=X.T)
+                # Make prediction using the processed data
+                prediction = model.predict(X)
+            except Exception as e:
+                return f"Error in foundation model preprocessing: {e}"
+        else:
+            # For classical models, use the data directly
+            prediction = model.predict(df)
+
+        # Convert prediction to human-readable output
+        result = prediction[0]
         if result == 1:
             return "Prediction: Positive for Diabetes"
         else:
@@ -152,11 +204,15 @@ def predict_diabetes(model_name, sex, race, ethnicity, marital_status, Prior_Mea
 
 # Define Gradio inputs
 inputs = [
-    gr.Dropdown(choices=get_available_models(), label="Select Model"),
+    gr.Dropdown(choices=["classical", "foundation"], label="Model Type"),
+    gr.Dropdown(choices=["diabetes", "24mths", "36mths", "48mths"], label="Time Period"),
     gr.Dropdown(choices=SEX_CHOICES, label="Sex"),
     gr.Dropdown(choices=RACE_CHOICES, label="Race"),
     gr.Dropdown(choices=ETHNICITY_CHOICES, label="Ethnicity"),
     gr.Dropdown(choices=MARITAL_STATUS_CHOICES, label="Marital Status"),
+    gr.Number(label="Age"), 
+    gr.Number(label="BMI"),
+    gr.Number(label="Smoking"), 
     gr.Number(label="Prior Mean Glu"),
     gr.Number(label="PT_ELX_GRP_1"), 
     gr.Number(label="PT_ELX_GRP_2"), 
@@ -202,8 +258,6 @@ inputs = [
     gr.Number(label="CAAA Drug"), 
     gr.Number(label="CCB Drug"),
     gr.Number(label="PAAAB Drug"), 
-    gr.Number(label="Age"), 
-    gr.Number(label="BMI"),
     gr.Number(label="Body Weight (kg)"), 
     gr.Number(label="SBP (Systolic Blood Pressure)"),
     gr.Number(label="DBP (Diastolic Blood Pressure)"), 
@@ -219,7 +273,6 @@ inputs = [
     gr.Number(label="Mean BUN"), 
     gr.Number(label="Mean AGAP"),
     gr.Number(label="Mean Protein"), 
-    gr.Number(label="Smoking"), 
     gr.Number(label="eGFR"),
     gr.Number(label="ED Visits"), 
     gr.Number(label="LOS (Length of Stay)"),