avoid non_null_labels key in loss kwargs

Author: aditya0by0

chebai/preprocessing/datasets/base.py

@@ -444,9 +444,12 @@ def _process_input_for_prediction(
         Returns:
             List[Dict[str, Any]]: Processed input data.
         """
+        # Add dummy labels because the collate function requires them.
+        # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`,
+        # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty.
         data = [
             self.reader.to_data(
-                {"id": f"smiles_{idx}", "features": smiles, "labels": None}
+                {"id": f"smiles_{idx}", "features": smiles, "labels": [1, 2]}
             )
             for idx, smiles in enumerate(smiles_list)
         ]

chebai/result/prediction.py

@@ -65,7 +65,7 @@ def predict_from_file(
 
         preds: torch.Tensor = self.predict_smiles(smiles=smiles_strings)
 
-        predictions_df = pd.DataFrame(torch.cat(preds).detach().cpu().numpy())
+        predictions_df = pd.DataFrame(preds.detach().cpu().numpy())
 
         def _add_class_columns(class_file_path: _PATH):
             with open(class_file_path, "r") as f: