Skip to content

Commit 4c9ebd5

Browse files
[PWGHF] Improve BDT model initialization in the charm–track femto producer (#14996)
1 parent eeb6d7b commit 4c9ebd5

File tree

1 file changed

+31
-45
lines changed

1 file changed

+31
-45
lines changed

PWGHF/HFC/TableProducer/producerCharmHadronsTrackFemtoDream.cxx

Lines changed: 31 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -302,56 +302,42 @@ struct HfProducerCharmHadronsTrackFemtoDream {
302302
bool useDstarMl = doprocessDataDstarToD0PiWithML || doprocessMcDstarToD0PiWithML;
303303

304304
if (applyMlMode == FillMlFromNewBDT) {
305-
if (useLcMl) {
306-
hfMlResponseLc.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
307-
hfMlResponseLc.cacheInputFeaturesIndices(namesInputFeatures);
308-
hfMlResponseLc.init();
309-
}
310-
if (useDplusMl) {
311-
hfMlResponseDplus.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
312-
hfMlResponseDplus.cacheInputFeaturesIndices(namesInputFeatures);
313-
hfMlResponseDplus.init();
314-
}
315-
if (useD0Ml) {
316-
hfMlResponseD0.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
317-
hfMlResponseD0.cacheInputFeaturesIndices(namesInputFeatures);
318-
hfMlResponseD0.init();
319-
}
320-
if (useDstarMl) {
321-
hfMlResponseDstar.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
322-
hfMlResponseDstar.cacheInputFeaturesIndices(namesInputFeatures);
323-
hfMlResponseDstar.init();
324-
}
325305

326-
if (loadModelsFromCCDB) {
327-
ccdbApi.init(ccdbUrl);
328-
if (useLcMl) {
329-
hfMlResponseLc.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
330-
}
331-
if (useDplusMl) {
332-
hfMlResponseDplus.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
333-
}
334-
if (useD0Ml) {
335-
hfMlResponseD0.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
336-
}
337-
if (useDstarMl) {
338-
hfMlResponseDstar.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
306+
auto setupFeatures = [&](auto& hfResponse, bool useMlFlag) {
307+
if (!useMlFlag) {
308+
return;
339309
}
310+
hfResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
311+
hfResponse.cacheInputFeaturesIndices(namesInputFeatures);
312+
};
340313

341-
} else {
342-
if (useLcMl) {
343-
hfMlResponseLc.setModelPathsLocal(onnxFileNames);
344-
}
345-
if (useDplusMl) {
346-
hfMlResponseDplus.setModelPathsLocal(onnxFileNames);
347-
}
348-
if (useD0Ml) {
349-
hfMlResponseD0.setModelPathsLocal(onnxFileNames);
314+
setupFeatures(hfMlResponseLc, useLcMl);
315+
setupFeatures(hfMlResponseDplus, useDplusMl);
316+
setupFeatures(hfMlResponseD0, useD0Ml);
317+
setupFeatures(hfMlResponseDstar, useDstarMl);
318+
319+
const bool useAnyMl = useLcMl || useDplusMl || useD0Ml || useDstarMl;
320+
if (loadModelsFromCCDB && useAnyMl) {
321+
ccdbApi.init(ccdbUrl);
322+
}
323+
324+
auto initModel = [&](auto& hfResponse, bool useMlFlag) {
325+
if (!useMlFlag) {
326+
return;
350327
}
351-
if (useDstarMl) {
352-
hfMlResponseDstar.setModelPathsLocal(onnxFileNames);
328+
329+
if (loadModelsFromCCDB) {
330+
hfResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
331+
} else {
332+
hfResponse.setModelPathsLocal(onnxFileNames);
353333
}
354-
}
334+
hfResponse.init();
335+
};
336+
337+
initModel(hfMlResponseLc, useLcMl);
338+
initModel(hfMlResponseDplus, useDplusMl);
339+
initModel(hfMlResponseD0, useD0Ml);
340+
initModel(hfMlResponseDstar, useDstarMl);
355341
}
356342
}
357343

0 commit comments

Comments
 (0)