[Code scan] Restore MolPredictHF multiclass and save_path behavior

[njzjz]

This issue is a result of a Codex global repository scan.

Summary

MolTrainHF does not persist multiclass_cnt, and MolPredictHF.predict treats every non-binary task as regression. It also saves only metric results when save_path is provided, instead of saving prediction CSV and metric JSON outputs like MolPredict.

Code references

unimol_tools/unimol_hf/trainer.py

Lines 130 to 136 in 4596596

	def _update_and_save_config(self):
	self.config["num_classes"] = self.data["num_classes"]
	self.config["target_cols"] = ",".join(self.data["target_cols"])
	self.config["split_method"] = f"{self.config['kfold']}fold_{self.config['split']}"
	os.makedirs(self.save_path, exist_ok=True)
	out_path = os.path.join(self.save_path, "config.yaml")
	self.yamlhandler.write_yaml(data=self.config, out_file_path=out_path)

unimol_tools/unimol_hf/trainer.py

Lines 152 to 188 in 4596596

	def predict(self, data, save_path=None, metrics="none"):
	if metrics and metrics != "none":
	self.config.metrics = metrics
	self.datahub = DataHub(
	data=data, is_train=False, save_path=self.load_model, **self.config
	)
	self.config.use_ddp = False
	self.trainer = Trainer(save_path=self.load_model, **self.config)
	self.model = HFNNModel(self.datahub.data, self.trainer, **self.config)
	self.model.evaluate(self.trainer, self.load_model)
	y_pred = self.model.cv["test_pred"]
	scalar = self.datahub.data["target_scaler"]
	if scalar is not None:
	y_pred = scalar.inverse_transform(y_pred)
	df = self.datahub.data["raw_data"].copy()
	predict_cols = ["predict_" + col for col in self.target_cols]
	if self.task in ["classification", "multilabel_classification"]:
	threshold = joblib.load(os.path.join(self.load_model, "threshold.dat"))
	prob_cols = ["prob_" + col for col in self.target_cols]
	df[prob_cols] = y_pred
	df[predict_cols] = (y_pred > threshold).astype(int)
	else:
	prob_cols = predict_cols
	df[predict_cols] = y_pred
	result_metrics = None
	if not (df[self.target_cols] == -1.0).all().all():
	result_metrics = self.trainer.metrics.cal_metric(
	df[self.target_cols].values, df[prob_cols].values
	)
	logger.info("final predict metrics score: \n{}".format(result_metrics))
	if save_path:
	os.makedirs(save_path, exist_ok=True)
	joblib.dump(result_metrics, os.path.join(save_path, "test_metric.result"))
	return y_pred, result_metrics

unimol_tools/unimol_tools/predict.py

Lines 87 to 121 in 4596596

	df = self.datahub.data['raw_data'].copy()
	predict_cols = ['predict_' + col for col in self.target_cols]
	if self.task == 'multiclass' and self.config.multiclass_cnt is not None:
	prob_cols = ['prob_' + str(i) for i in range(self.config.multiclass_cnt)]
	df[prob_cols] = y_pred
	df[predict_cols] = np.argmax(y_pred, axis=1).reshape(-1, 1)
	elif self.task in ['classification', 'multilabel_classification']:
	threshold = joblib.load(
	open(os.path.join(self.load_model, 'threshold.dat'), "rb")
	)
	prob_cols = ['prob_' + col for col in self.target_cols]
	df[prob_cols] = y_pred
	df[predict_cols] = (y_pred > threshold).astype(int)
	else:
	prob_cols = predict_cols
	df[predict_cols] = y_pred
	if self.save_path:
	os.makedirs(self.save_path, exist_ok=True)
	if not (df[self.target_cols] == -1.0).all().all():
	metrics = self.trainer.metrics.cal_metric(
	df[self.target_cols].values, df[prob_cols].values
	)
	logger.info("final predict metrics score: \n{}".format(metrics))
	if self.save_path:
	joblib.dump(metrics, os.path.join(self.save_path, 'test_metric.result'))
	with open(os.path.join(self.save_path, 'test_metric.json'), 'w') as f:
	json.dump(metrics, f)
	else:
	df.drop(self.target_cols, axis=1, inplace=True)
	if self.save_path:
	prefix = (
	data.split('/')[-1].split('.')[0] if isinstance(data, str) else 'test'
	)
	self.save_predict(df, self.save_path, prefix)
	logger.info("pipeline finish!")

Impact

HF multiclass predictions can be written into the wrong columns and scored with the wrong target/prediction shape. Users also lose the prediction CSV output that the regular MolPredict path produces.

Suggested fix

Mirror MolPredict behavior: persist multiclass_cnt during training, add a multiclass prediction branch with probability columns and argmax labels, set prediction save paths, and save prediction CSV plus metric JSON/result files.