Fixed the test after rebase

Signed-off-by: Meet Patel <[email protected]>

Author: quic-meetkuma

QEfficient/cloud/finetune.py

@@ -294,7 +294,7 @@ def setup_dataloaders(
     return train_dataloader, eval_dataloader, longest_seq_length
 
 
-def main(peft_config_file=None, **kwargs) -> None:
+def main(peft_config_file: str = None, **kwargs) -> None:
     """
     Fine-tune a model on QAIC hardware with configurable training and LoRA parameters.
 

QEfficient/finetune/utils/config_utils.py

@@ -84,7 +84,7 @@ def generate_peft_config(train_config: TrainConfig, peft_config_file: str = None
         if train_config.peft_method not in config_map:
             raise RuntimeError(f"Peft config not found: {train_config.peft_method}")
 
-        config_cls, peft_config_cls = config_map[train_config.peft_method]()
+        config_cls, peft_config_cls = config_map[train_config.peft_method]
         if config_cls is None:
             params = kwargs
         else:

QEfficient/finetune/utils/train_utils.py

@@ -458,7 +458,7 @@ def evaluation_helper(model, train_config, eval_dataloader, device):
     # Print evaluation metrics
     print(f" {eval_metric.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
 
-    return eval_metric, eval_epoch_loss, val_step_loss, val_step_metric
+    return eval_epoch_loss, eval_metric, val_step_loss, val_step_metric
 
 
 def get_longest_seq_length(data: List[Dict]) -> Tuple[int, int]:

tests/finetune/test_finetune.py

@@ -66,11 +66,10 @@ def test_finetune(
     }
 
     results = finetune(**kwargs)
-
-    assert np.allclose(results["avg_train_metric"], 1.002326, atol=1e-5), "Train metric is not matching."
     assert np.allclose(results["avg_train_loss"], 0.00232327, atol=1e-5), "Train loss is not matching."
-    assert np.allclose(results["avg_eval_metric"], 1.0193923, atol=1e-5), "Eval metric is not matching."
-    assert np.allclose(results["avg_eval_loss"], 0.0192067, atol=1e-5), "Eval loss is not matching."
+    assert np.allclose(results["avg_train_metric"], 1.002326, atol=1e-5), "Train metric is not matching."
+    assert np.allclose(results["avg_eval_loss"], 0.0206124, atol=1e-5), "Eval loss is not matching."
+    assert np.allclose(results["avg_eval_metric"], 1.020826, atol=1e-5), "Eval metric is not matching."
     assert results["avg_epoch_time"] < 60, "Training should complete within 60 seconds."
 
     train_config_spy.assert_called_once()
@@ -86,8 +85,8 @@ def test_finetune(
     assert get_preprocessed_dataset_spy.call_count == 2
 
     args, kwargs = train_spy.call_args
-    train_dataloader = args[1]
-    eval_dataloader = args[2]
+    train_dataloader = args[2]
+    eval_dataloader = args[3]
     optimizer = args[4]
 
     batch = next(iter(train_dataloader))