Made changes to optimize hyperparams for ablation study

experiments/train_optuna.py

@@ -1,26 +1,29 @@
 import os
 import sys
+import torch 
+import optuna
+import yaml
+import json
+import wandb 
+import argparse
+import matplotlib.pyplot as plt 
+import seaborn as sns 
+
+from tqdm import tqdm
+from torch.utils.data import DataLoader, Subset 
+from torch.nn import CrossEntropyLoss
+from torch.optim import Adam 
+from sklearn.model_selection import StratifiedShuffleSplit
+from sklearn.metrics import accuracy_score, f1_score, classification_report, confusion_matrix
 
+
+# Setup base path
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 # Automatically add Project root to python import path
 base_dir = os.path.dirname(os.path.dirname(__file__))
 if base_dir not in sys.path:
     sys.path.append(base_dir)
 
-import torch 
-import optuna 
-from torch.utils.data import DataLoader, Subset 
-from torch.nn import CrossEntropyLoss
-from torch.optim import Adam 
-from sklearn.model_selection import StratifiedShuffleSplit
-from sklearn.metrics import accuracy_score, f1_score, classification_report, confusion_matrix
-from tqdm import tqdm
-import matplotlib.pyplot as plt 
-import seaborn as sns 
-import wandb 
-import json
-import yaml
-import argparse
 
 from src.triage_dataset import TriageDataset
 from src.multimodal_model import MediLLMModel
@@ -36,12 +39,13 @@ def stratified_split(dataset, val_ratio=0.2, seed=42, label_column="triage_level
 
 def objective(trial):
     wandb.init(
-        project="mediLLM-v2",
-        name=f"trial-{trial.number}-v4-{wandb.util.generate_id()}",
+        project=f"mediLLM-tune-{mode}",
+        name=f"{mode}-trial-{trial.number}-v5-{wandb.util.generate_id()}",
         group="SoftLabelTrials",
         config={
             "dataset_version": "softlabels",
-            "dataset_size": 900
+            "dataset_size": 900,
+            "mode": mode
         }
     )
 
@@ -51,7 +55,7 @@ def objective(trial):
     hidden_dim = trial.suggest_categorical("hidden_dim", [128, 256, 512])
     batch_size = trial.suggest_categorical("bs", [4, 8, 16])
 
-    model = MediLLMModel(dropout=dropout, hidden_dim=hidden_dim).to(device)
+    model = MediLLMModel(dropout=dropout, hidden_dim=hidden_dim, mode=mode).to(device)
     wandb.watch(model)
 
     dataset = TriageDataset(os.path.join(base_dir, "data", "emr_records.csv"))
@@ -65,132 +69,155 @@ def objective(trial):
 
     for epoch in range(2):
         model.train()
-        loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/2", leave=False)
+        loop = tqdm(train_loader, desc=f"[{mode}] Epoch {epoch+1}/2", leave=False)
         for batch in loop:
-            input_ids = batch["input_ids"].to(device)
-            attention_mask = batch["attention_mask"].to(device)
-            images = batch["image"].to(device)
+            input_ids = batch.get("input_ids", None)
+            attention_mask = batch.get("attention_mask", None)
+            images = batch.get("image", None)
             labels = batch["label"].to(device)
 
+            if input_ids is not None:
+                input_ids = input_ids.to(device)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+            if images is not None:
+                images = images.to(device)
+
             optimizer.zero_grad()
             outputs = model(input_ids=input_ids, attention_mask=attention_mask, image=images)
             loss = criterion(outputs, labels)
             loss.backward()
             optimizer.step()
-
             loop.set_postfix(loss=loss.item())
 
     # Validation
     model.eval()
     all_preds, all_labels = [], []
     with torch.no_grad():
-        for batch in tqdm(val_loader, desc="Validating", leave=False):
-            input_ids = batch["input_ids"].to(device)
-            attention_mask = batch["attention_mask"].to(device)
-            images = batch["image"].to(device)
+        for batch in tqdm(val_loader, desc=f"[{mode}] Validating", leave=False):
+            input_ids = batch.get("input_ids", None)
+            attention_mask = batch.get("attention_mask", None)
+            images = batch.get("image", None)
             labels = batch["label"].to(device)
 
+            if input_ids is not None:
+                input_ids = input_ids.to(device)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+            if images is not None:
+                images = images.to(device)
+
             outputs = model(input_ids=input_ids, attention_mask=attention_mask, image=images)
             preds = torch.argmax(outputs, dim=1).cpu().numpy()
             all_preds.extend(preds)
             all_labels.extend(labels.cpu().numpy())
 
     f1 = f1_score(all_labels, all_preds, average="weighted")
-    print(f"\n[Trial {trial.number}] Classification Report:")
-    print(classification_report(all_labels, all_preds, target_names=["low", "medium", "high"]))
+    acc = accuracy_score(all_labels, all_preds)
 
+    # Log to W&B and Optuna
+    wandb.log({
+        "val_f1_score": f1,
+        "val_accuracy": acc,
+        "lr": lr,
+        "dropout": dropout,
+        "hidden_dim": hidden_dim,
+        "batch_size": batch_size
+    })
+
+    # Confusion Matrix
     cm = confusion_matrix(all_labels, all_preds)
     plt.figure(figsize=(6, 5))
     sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                 xticklabels=["low", "medium", "high"],
                 yticklabels=["low", "medium", "high"])
-    plt.title(f"Confusion Matrix - Trial {trial.number}")
+    plt.title(f"Confusion Matrix - {mode} Trial {trial.number}")
     plt.xlabel("Predicted")
     plt.ylabel("True")
-    wandb.log({f"confusion_matrix/trial_{trial.number}": wandb.Image(plt)})
+    wandb.log({f"{mode}_confusion_matrix/trial_{trial.number}": wandb.Image(plt)})
     plt.close()
     
-    # Log to W&B and Optuna
-    wandb.log({
-        "f1_score": f1,
-        "accuracy": accuracy_score(all_labels, all_preds),
-        "lr": lr,
-        "dropout": dropout,
-        "hidden_dim": hidden_dim,
-        "batch_size": batch_size
-    })
     return f1
 
 def get_args():
     parser = argparse.ArgumentParser(description="Run Optuna hyperparameter search")
     parser.add_argument("--n_trials", type=int, default=10, help="Number of Optuna trials to run")
+    parser.add_argument("--mode", type=str, choices=["text", "image", "multimodal"], required=True, help="Input mode")
     return parser.parse_args()
 
 if __name__=="__main__":
     args = get_args()
+    mode = args.mode 
 
     study = optuna.create_study(
-        study_name="mediLLM_v2",
+        study_name=f"mediLLM_{mode}_optuna",
         direction="maximize"
     )
-    with tqdm(total=args.n_trials, desc="Optuna Trials") as pbar:
+    with tqdm(total=args.n_trials, desc=f"Optuna Trials [{mode}]") as pbar:
         def wrapped_objective(trial):
             try:
-                result = objective(trial)
-                return result
+                return objective(trial, mode)
             finally:
                 wandb.finish()
                 pbar.update(1)
     
         study.optimize(wrapped_objective, n_trials=args.n_trials)
 
-    print("Best F1 score achieved:", study.best_value)
-    print("Best hyperparameters:", study.best_params)
+    print(f"✅ Best F1 score for {mode}: {study.best_value}")
+    print(f"✅ Best hyperparameters: {study.best_params}")
+
+    # Save best hyperparameters to JSON per mode
+    json_path = os.path.join(base_dir, "assets", "best_hyperparams.json")
+    os.makedirs(os.path.dirname(json_path), exist_ok=True)
 
-    # Save as JSON
-    assets_dir = os.path.join(base_dir, "assets")
+    best_params_entry = {
+        "lr": float(study.best_params["lr"]),
+        "dropout": float(study.best_params["dropout"]),
+        "hidden_dim": int(study.best_params["hidden_dim"]),
+        "batch_size": int(study.best_params["bs"]),
+        "epochs": 5
+    }
 
-    # Make sure assets directory exists in the root
-    os.makedirs(assets_dir, exist_ok=True)
+    # Load existing or start new
+    if os.path.exists(json_path):
+        with open(json_path, "r") as f:
+            best_params_all = json.load(f)
 
-    # Save the best hyperparameters
-    with open(os.path.join(assets_dir, "best_hyperparams.json"), "w") as f:
-        json.dump(study.best_params, f, indent=4)
+    else:
+        best_params_all = {}
+
+    best_params_all[mode] = best_params_entry
+
+    # Write back
+    with open(json_path, "w") as f:
+        json.dump(best_params_all, f, indent=4)
+
+    print(f"✅ Saved best hyperparameters for [{mode}] to best_hyperparams.json")
 
     # Export to config.yaml
-    config_dir = os.path.join(base_dir, "config")
-    config_path = os.path.join(config_dir, "config.yaml")
+    config_path = os.path.join(base_dir, "config", "config.yaml")
     
     # Make sure config directory exists in the root
-    os.makedirs(config_dir, exist_ok=True)
-
-    # If the config file doesn't exist, create a default one
-    if not os.path.exists(config_path):
-        with open(config_path, "w") as f:
-            f.write(
-                "model:\n"
-                "  dropout: 0.3\n"
-                "  hidden_dim: 256\n\n"
-                "train:\n"
-                "  lr: 2e-5\n"
-                "  batch_size: 8\n"
-                "  epochs: 5\n\n"
-                "wandb:\n"
-                " project: medi-llm-final\n"
-            )
+    os.makedirs(os.path.dirname(config_path), exist_ok=True)
 
-    # Export to config.yaml
-    with open(config_path, "r") as f:
-        cfg = yaml.safe_load(f)
+    config = {}
+    if os.path.exists(config_path):
+        with open(config_path, "r") as f:
+            config = yaml.safe_load(f) or {}
 
-    cfg["model"]["dropout"] = float(study.best_params["dropout"])
-    cfg["model"]["hidden_dim"] = int(study.best_params["hidden_dim"])
-    cfg["train"]["lr"] = float(study.best_params["lr"])
-    cfg["train"]["batch_size"] = int(study.best_params["bs"])
+    config[mode] = {
+        "lr": float(study.best_params["lr"]),
+        "dropout": float(study.best_params["dropout"]),
+        "hidden_dim": int(study.best_params["hidden_dim"]),
+        "batch_size": int(study.best_params["bs"]),
+        "epochs": 5
+    }
 
-    # Save updated config
+    # Export to config.yaml
     with open(config_path, "w") as f:
-        yaml.dump(cfg, f, default_flow_style=False)
+        yaml.dump(config, f, sort_keys=False)
+
+    print(f"✅ Best hyperparameters for [{mode}] saved in config.yaml")
     
 
     

src/train.py

@@ -1,50 +1,69 @@
 import torch # PyTorch core utility for model training 
 import os
 import sys
+import yaml
+import argparse
+import matplotlib.pyplot as plt # for plotting
+
+from tqdm import tqdm # loading bar for loops
+from torch.utils.data import DataLoader, Subset # Dataloader to batch and feed data to model, random split to split dataset into train and validation sets
+from torch.nn import CrossEntropyLoss # PyTorch core utility for model training
+from torch.optim import Adam # PyTorch core utility for model training, Adam is the Optimizer a gradient descent model
+from sklearn.metrics import accuracy_score, f1_score # Evaluation metrics
+from sklearn.model_selection import StratifiedShuffleSplit
+
 
+# Setup base path
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 # Automatically add Project root to python import path
 base_dir = os.path.dirname(os.path.dirname(__file__))
 if base_dir not in sys.path:
     sys.path.append(base_dir)
 
-import argparse
-import yaml
-from torch.utils.data import DataLoader, Subset # Dataloader to batch and feed data to model, random split to split dataset into train and validation sets
-from torch.nn import CrossEntropyLoss # PyTorch core utility for model training
-from torch.optim import Adam # PyTorch core utility for model training, Adam is the Optimizer a gradient descent model
-from sklearn.metrics import accuracy_score, f1_score # Evaluation metrics
-from sklearn.model_selection import StratifiedShuffleSplit
-from tqdm import tqdm # loading bar for loops
-import matplotlib.pyplot as plt # for plotting
+
 from src.triage_dataset import TriageDataset # Dataset Class
 from src.multimodal_model import MediLLMModel # Mutlimodal Model
 
-def load_config():
-    config_dir = os.path.join(base_dir, "config")
-    config_path = os.path.join(config_dir, "config.yaml")
+def load_config(mode):
+    config_path = os.path.join(base_dir, "config", "config.yaml")
+    os.makedirs(os.path.dirname(config_path), exist_ok=True)
 
-    # Make sure config directory exists in the root
-    os.makedirs(config_dir, exist_ok=True)
-
-     # If the config file doesn't exist, create a default one
+    # If the config file doesn't exist, create it defaults for all modes
     if not os.path.exists(config_path):
+        default_config = {
+            "text": {
+                "lr": 2e-5,
+                "dropout": 0.3,
+                "hidden_dim": 256,
+                "batch_size": 8,
+                "epochs": 5
+            },
+            "image": {
+                "lr": 2e-5,
+                "dropout": 0.3,
+                "hidden_dim": 256,
+                "batch_size": 8,
+                "epochs": 5
+            },
+            "multimodal": {
+                "lr": 2e-5,
+                "dropout": 0.3,
+                "hidden_dim": 256,
+                "batch_size": 8,
+                "epochs": 5
+            }
+        }
         with open(config_path, "w") as f:
-            f.write(
-                "model:\n"
-                "  dropout: 0.3\n"
-                "  hidden_dim: 256\n\n"
-                "train:\n"
-                "  lr: 2e-5\n"
-                "  batch_size: 8\n"
-                "  epochs: 5\n\n"
-                "wandb:\n"
-                " project: medi-llm-final\n"
-            )
+            yaml.dump(default_config, f)
     
     # otherwise export to yaml
     with open(config_path, "r") as f:
-        return yaml.safe_load(f)
+        config = yaml.safe_load(f)
+
+    if mode not in config:
+        raise ValueError(f"No config found for mode '{mode}' in config.yaml")
+    
+    return config[mode]
 
 def stratified_split(dataset, val_ratio=0.2, seed=42):
     labels = [dataset.df.iloc[i]["triage_level"] for i in range(len(dataset))]
@@ -53,33 +72,33 @@ def stratified_split(dataset, val_ratio=0.2, seed=42):
     return Subset(dataset, tran_idx), Subset(dataset, val_idx)
 
 def train_model(mode="multimodal"): # Function to instantiate model and data, train, validate, plot results and save the model
-    config = load_config()
+    cfg = load_config(mode)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Use GPU if available or else use CPU
     
-    dataset_dir = os.path.join(base_dir, "data", "emr_records_softlabels.csv")
+    dataset_dir = os.path.join(base_dir, "data", "emr_records.csv")
     dataset = TriageDataset(
         csv_file=dataset_dir,
         mode=mode
     )
 
     model = MediLLMModel(
-        dropout=config["model"]["dropout"],
-        hidden_dim=config["model"]["hidden_dim"],
+        dropout=cfg["dropout"],
+        hidden_dim=cfg["hidden_dim"],
         mode = mode
     ).to(device) # moves the model to selected device
  
     train_set, val_set = stratified_split(dataset)
-    batch_size = config["train"]["batch_size"]
+    batch_size = cfg["batch_size"]
 
     train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True) # Create data in batches to the model
     val_loader = DataLoader(val_set, batch_size=batch_size)
 
     criterion = CrossEntropyLoss() # Calculate difference between model prediction and true labels
-    optimizer = Adam(model.parameters(), lr=config["train"]["lr"]) # Adaptive learning rate optimizer for fast-converging
+    optimizer = Adam(model.parameters(), lr=cfg["lr"]) # Adaptive learning rate optimizer for fast-converging
 
     train_acc, val_acc = [], [] # Lists to store accuracy per epoch for plotting
 
-    for epoch in range(config["train"]["epochs"]):
+    for epoch in range(cfg["epochs"]):
         model.train() # Activate training the model, enable dropout
         all_preds, all_labels = [], []
 
@@ -152,11 +171,11 @@ def train_model(mode="multimodal"): # Function to instantiate model and data, tr
         print(f"Val Accuracy: {val_acc_epoch:.4f}, F1 Score: {val_f1:.4f}")
 
     # Save model
-    model_path = os.path.join(base_dir, f"medi_llm_model_softlabels{mode}.pth")
+    model_path = os.path.join(base_dir, f"medi_llm_model_{mode}.pth")
     torch.save(model.state_dict(), model_path) # Saves the model weights only not total architecture to reuse later
 
     # Plot accuracy
-    plot_path = os.path.join(base_dir, "assets", f"model_training_curve_softlabels{mode}.png")
+    plot_path = os.path.join(base_dir, "assets", f"model_training_curve_{mode}.png")
     plt.plot(train_acc, label="Train Acc")
     plt.plot(val_acc, label="Val Acc")
     plt.legend()