diff --git a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py
index 396fcefd84a79..f349808c487f7 100644
--- a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py
+++ b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py
@@ -5,6 +5,8 @@
from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton
from Configuration.ProcessModifiers.deepMETSonicTriton_cff import deepMETSonicTriton
from Configuration.ProcessModifiers.deepTauSonicTriton_cff import deepTauSonicTriton
+from Configuration.ProcessModifiers.particleTransformerAK4SonicTriton_cff import particleTransformerAK4SonicTriton
# collect all SonicTriton-related process modifiers here
-allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton)
+allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton)
+
diff --git a/Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py b/Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py
new file mode 100644
index 0000000000000..22fff0fb9ea6c
--- /dev/null
+++ b/Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py
@@ -0,0 +1,3 @@
+import FWCore.ParameterSet.Config as cms
+
+particleTransformerAK4SonicTriton = cms.Modifier()
diff --git a/RecoBTag/ONNXRuntime/BuildFile.xml b/RecoBTag/ONNXRuntime/BuildFile.xml
index c9b2daa7fe264..45e88927184fa 100644
--- a/RecoBTag/ONNXRuntime/BuildFile.xml
+++ b/RecoBTag/ONNXRuntime/BuildFile.xml
@@ -1,3 +1,4 @@
+
diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h
index e528b4e92f74a..98faeda8403db 100644
--- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h
+++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h
@@ -1,6 +1,7 @@
#ifndef RecoBTag_ONNXRuntime_tensor_configs_h
#define RecoBTag_ONNXRuntime_tensor_configs_h
+#include
namespace deepflavour {
constexpr unsigned n_features_global = 15;
@@ -28,4 +29,41 @@ namespace deepvertex {
} // namespace deepvertex
+namespace parT {
+
+ enum InputFeatures {
+ kBegin = 0,
+ kChargedCandidates = kBegin,
+ kNeutralCandidates = 1,
+ kVertices = 2,
+ kChargedCandidates4Vec = 3,
+ kNeutralCandidates4Vec = 4,
+ kVertices4Vec = 5,
+ kEnd = 6
+ };
+
+ inline constexpr unsigned n_cpf_accept = 25;
+ inline constexpr unsigned n_npf_accept = 25;
+ inline constexpr unsigned n_sv_accept = 5;
+
+ constexpr std::array N_InputFeatures{{
+ 16, // kChargedCandidates
+ 8, // kNeutralCandidates
+ 14, // kVertices
+ 4, // kChargedCandidates4Vec
+ 4, // kNeutralCandidates4Vec
+ 4, // kVertices4Vec
+ }};
+
+ constexpr std::array N_AcceptedFeatures{{
+ n_cpf_accept, // kChargedCandidates
+ n_npf_accept, // kNeutralCandidates
+ n_sv_accept, // kVertices
+ n_cpf_accept, // kChargedCandidates4Vec
+ n_npf_accept, // kNeutralCandidates4Vec
+ n_sv_accept, // kVertices4Vec
+ }};
+
+} // namespace parT
+
#endif
diff --git a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h
index bb1a25a2d0f0b..db3facc2338d1 100644
--- a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h
+++ b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h
@@ -2,6 +2,9 @@
#define RecoBTag_ONNXRuntime_tensor_fillers_h
#include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h"
+#include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
+#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
namespace btagbtvdeep {
@@ -19,6 +22,53 @@ namespace btagbtvdeep {
void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features);
+ std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features,
+ parT::InputFeatures ifeature);
+
+ std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features,
+ parT::InputFeatures ifeature);
+
+ std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature);
+
+ template
+ void parT_tensor_filler(cms::Ort::FloatArrays& data,
+ const parT::InputFeatures ifeature,
+ const std::vector& features,
+ const unsigned int max_n,
+ const float*& start,
+ unsigned offset) {
+ float* ptr = nullptr;
+ for (std::size_t n = 0; n < max_n; n++) {
+ const auto& f = features.at(n);
+ ptr = &data[ifeature][offset + n * parT::N_InputFeatures.at(ifeature)];
+ start = ptr;
+ const std::vector& inputs = inputs_parT(f, ifeature);
+ for (unsigned int i = 0; i < inputs.size(); i++) {
+ *ptr = inputs[i];
+ ++ptr;
+ }
+ if (inputs.size() > 0)
+ --ptr;
+ assert(start + parT::N_InputFeatures.at(ifeature) - 1 == ptr);
+ }
+ }
+
+ template
+ void parT_tensor_filler(std::vector& vdata,
+ const parT::InputFeatures ifeature,
+ const std::vector& features,
+ const unsigned int target_n) {
+ unsigned int n =
+ std::clamp((unsigned int)features.size(), (unsigned int)0, (unsigned int)parT::N_AcceptedFeatures.at(ifeature));
+ for (unsigned int count = 0; count < n; count++) {
+ const std::vector& inputs = inputs_parT(features.at(count), ifeature);
+ vdata.insert(vdata.end(), inputs.begin(), inputs.end());
+ }
+ unsigned int n_features = parT::N_InputFeatures.at(ifeature);
+ if (n < target_n)
+ vdata.insert(vdata.end(), (target_n - n) * n_features, 0); // Add 0 to unfilled part as padding value
+ }
+
} // namespace btagbtvdeep
#endif
diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc
index b27a9b397e520..ea5fc6b0020f4 100644
--- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc
+++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc
@@ -16,6 +16,9 @@
#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
+
using namespace cms::Ort;
class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer> {
@@ -41,24 +44,9 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer
std::vector flav_names_;
std::vector input_names_;
std::vector output_names_;
-
- enum InputIndexes {
- kChargedCandidates = 0,
- kNeutralCandidates = 1,
- kVertices = 2,
- kChargedCandidates4Vec = 3,
- kNeutralCandidates4Vec = 4,
- kVertices4Vec = 5
- };
- unsigned n_cpf_;
- constexpr static unsigned n_features_cpf_ = 16;
- constexpr static unsigned n_pairwise_features_cpf_ = 4;
- unsigned n_npf_;
- constexpr static unsigned n_features_npf_ = 8;
- constexpr static unsigned n_pairwise_features_npf_ = 4;
- unsigned n_sv_;
- constexpr static unsigned n_features_sv_ = 14;
- constexpr static unsigned n_pairwise_features_sv_ = 4;
+ unsigned int n_cpf_;
+ unsigned int n_npf_;
+ unsigned int n_sv_;
std::vector input_sizes_;
std::vector> input_shapes_; // shapes of each input group (-1 for dynamic axis)
@@ -84,7 +72,7 @@ void ParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Configurat
desc.add("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
desc.add>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
desc.add("model_path",
- edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/RobustParTAK4.onnx"));
+ edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/modelfile/model.onnx"));
desc.add>("output_names", {"softmax"});
desc.add>(
"flav_names", std::vector{"probb", "probbb", "problepb", "probc", "probuds", "probg"});
@@ -124,12 +112,12 @@ void ParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEvent, cons
get_input_sizes(taginfo);
// run prediction with dynamic batch size per event
- input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)n_features_cpf_},
- {(int64_t)1, (int64_t)n_npf_, (int64_t)n_features_npf_},
- {(int64_t)1, (int64_t)n_sv_, (int64_t)n_features_sv_},
- {(int64_t)1, (int64_t)n_cpf_, (int64_t)n_pairwise_features_cpf_},
- {(int64_t)1, (int64_t)n_npf_, (int64_t)n_pairwise_features_npf_},
- {(int64_t)1, (int64_t)n_sv_, (int64_t)n_pairwise_features_sv_}};
+ input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates)},
+ {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates)},
+ {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices)},
+ {(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates4Vec)},
+ {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec)},
+ {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices4Vec)}};
outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0];
assert(outputs.size() == flav_names_.size());
@@ -151,24 +139,17 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
const reco::FeaturesTagInfo taginfo) {
const auto& features = taginfo.features();
- unsigned int n_cpf = features.c_pf_features.size();
- unsigned int n_npf = features.n_pf_features.size();
- unsigned int n_vtx = features.sv_features.size();
+ n_cpf_ = std::clamp((unsigned int)features.c_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_cpf_accept);
+ n_npf_ = std::clamp((unsigned int)features.n_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_npf_accept);
+ n_sv_ = std::clamp((unsigned int)features.sv_features.size(), (unsigned int)1, (unsigned int)parT::n_sv_accept);
- n_cpf_ = std::max((unsigned int)1, n_cpf);
- n_npf_ = std::max((unsigned int)1, n_npf);
- n_sv_ = std::max((unsigned int)1, n_vtx);
-
- n_cpf_ = std::min((unsigned int)25, n_cpf_);
- n_npf_ = std::min((unsigned int)25, n_npf_);
- n_sv_ = std::min((unsigned int)5, n_sv_);
input_sizes_ = {
- n_cpf_ * n_features_cpf_,
- n_npf_ * n_features_npf_,
- n_sv_ * n_features_sv_,
- n_cpf_ * n_pairwise_features_cpf_,
- n_npf_ * n_pairwise_features_npf_,
- n_sv_ * n_pairwise_features_sv_,
+ n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates),
+ n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates),
+ n_sv_ * parT::N_InputFeatures.at(parT::kVertices),
+ n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates4Vec),
+ n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec),
+ n_sv_ * parT::N_InputFeatures.at(parT::kVertices4Vec),
};
// init data storage
data_.clear();
@@ -180,116 +161,26 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
}
void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::ParticleTransformerAK4Features features) {
- float* ptr = nullptr;
+ //float* ptr = nullptr;
const float* start = nullptr;
unsigned offset = 0;
- // c_pf candidates
auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
- for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) {
- const auto& c_pf_features = features.c_pf_features.at(c_pf_n);
- ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_];
- start = ptr;
- *ptr = c_pf_features.btagPf_trackEtaRel;
- *(++ptr) = c_pf_features.btagPf_trackPtRel;
- *(++ptr) = c_pf_features.btagPf_trackPPar;
- *(++ptr) = c_pf_features.btagPf_trackDeltaR;
- *(++ptr) = c_pf_features.btagPf_trackPParRatio;
- *(++ptr) = c_pf_features.btagPf_trackSip2dVal;
- *(++ptr) = c_pf_features.btagPf_trackSip2dSig;
- *(++ptr) = c_pf_features.btagPf_trackSip3dVal;
- *(++ptr) = c_pf_features.btagPf_trackSip3dSig;
- *(++ptr) = c_pf_features.btagPf_trackJetDistVal;
- *(++ptr) = c_pf_features.ptrel;
- *(++ptr) = c_pf_features.drminsv;
- *(++ptr) = c_pf_features.vtx_ass;
- *(++ptr) = c_pf_features.puppiw;
- *(++ptr) = c_pf_features.chi2;
- *(++ptr) = c_pf_features.quality;
- assert(start + n_features_cpf_ - 1 == ptr);
- }
-
- // n_pf candidates
auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
- for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) {
- const auto& n_pf_features = features.n_pf_features.at(n_pf_n);
- ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_];
- start = ptr;
- *ptr = n_pf_features.ptrel;
- *(++ptr) = n_pf_features.etarel;
- *(++ptr) = n_pf_features.phirel;
- *(++ptr) = n_pf_features.deltaR;
- *(++ptr) = n_pf_features.isGamma;
- *(++ptr) = n_pf_features.hadFrac;
- *(++ptr) = n_pf_features.drminsv;
- *(++ptr) = n_pf_features.puppiw;
- assert(start + n_features_npf_ - 1 == ptr);
- }
-
- // sv candidates
auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);
- for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) {
- const auto& sv_features = features.sv_features.at(sv_n);
- ptr = &data_[kVertices][offset + sv_n * n_features_sv_];
- start = ptr;
- *ptr = sv_features.pt;
- *(++ptr) = sv_features.deltaR;
- *(++ptr) = sv_features.mass;
- *(++ptr) = sv_features.etarel;
- *(++ptr) = sv_features.phirel;
- *(++ptr) = sv_features.ntracks;
- *(++ptr) = sv_features.chi2;
- *(++ptr) = sv_features.normchi2;
- *(++ptr) = sv_features.dxy;
- *(++ptr) = sv_features.dxysig;
- *(++ptr) = sv_features.d3d;
- *(++ptr) = sv_features.d3dsig;
- *(++ptr) = sv_features.costhetasvpv;
- *(++ptr) = sv_features.enratio;
- assert(start + n_features_sv_ - 1 == ptr);
- }
+ // c_pf candidates
+ parT_tensor_filler(data_, parT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset);
+ // n_pf candidates
+ parT_tensor_filler(data_, parT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset);
+ // sv candidates
+ parT_tensor_filler(data_, parT::kVertices, features.sv_features, max_sv_n, start, offset);
// cpf pairwise features (4-vectors)
- auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
- for (std::size_t cpf_n = 0; cpf_n < max_cpf_n; cpf_n++) {
- const auto& cpf_pairwise_features = features.c_pf_features.at(cpf_n);
- ptr = &data_[kChargedCandidates4Vec][offset + cpf_n * n_pairwise_features_cpf_];
- start = ptr;
- *ptr = cpf_pairwise_features.px;
- *(++ptr) = cpf_pairwise_features.py;
- *(++ptr) = cpf_pairwise_features.pz;
- *(++ptr) = cpf_pairwise_features.e;
-
- assert(start + n_pairwise_features_cpf_ - 1 == ptr);
- }
-
+ parT_tensor_filler(data_, parT::kChargedCandidates4Vec, features.c_pf_features, max_c_pf_n, start, offset);
// npf pairwise features (4-vectors)
- auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
- for (std::size_t npf_n = 0; npf_n < max_npf_n; npf_n++) {
- const auto& npf_pairwise_features = features.n_pf_features.at(npf_n);
- ptr = &data_[kNeutralCandidates4Vec][offset + npf_n * n_pairwise_features_npf_];
- start = ptr;
- *ptr = npf_pairwise_features.px;
- *(++ptr) = npf_pairwise_features.py;
- *(++ptr) = npf_pairwise_features.pz;
- *(++ptr) = npf_pairwise_features.e;
-
- assert(start + n_pairwise_features_npf_ - 1 == ptr);
- }
-
+ parT_tensor_filler(data_, parT::kNeutralCandidates4Vec, features.n_pf_features, max_n_pf_n, start, offset);
// sv pairwise features (4-vectors)
- auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_);
- for (std::size_t sv_N = 0; sv_N < max_sv_N; sv_N++) {
- const auto& sv_pairwise_features = features.sv_features.at(sv_N);
- ptr = &data_[kVertices4Vec][offset + sv_N * n_pairwise_features_sv_];
- start = ptr;
- *ptr = sv_pairwise_features.px;
- *(++ptr) = sv_pairwise_features.py;
- *(++ptr) = sv_pairwise_features.pz;
- *(++ptr) = sv_pairwise_features.e;
-
- assert(start + n_pairwise_features_sv_ - 1 == ptr);
- }
+ parT_tensor_filler(data_, parT::kVertices4Vec, features.sv_features, max_sv_n, start, offset);
}
//define this as a plug-in
diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc
new file mode 100644
index 0000000000000..2a33e976e61f0
--- /dev/null
+++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc
@@ -0,0 +1,191 @@
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+#include "FWCore/Framework/interface/makeRefToBaseProdFrom.h"
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/Utilities/interface/StreamID.h"
+
+#include "DataFormats/BTauReco/interface/JetTag.h"
+
+#include "DataFormats/BTauReco/interface/ParticleTransformerAK4TagInfo.h"
+#include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
+
+#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
+#include "HeterogeneousCore/SonicTriton/interface/TritonData.h"
+
+#include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
+
+class ParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> {
+public:
+ explicit ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet &);
+ ~ParticleTransformerAK4SonicJetTagsProducer() override;
+
+ void acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) override;
+
+ void produce(edm::Event &iEvent, edm::EventSetup const &iSetup, Output const &iOutput) override;
+ static void fillDescriptions(edm::ConfigurationDescriptions &);
+
+private:
+ typedef std::vector TagInfoCollection;
+ typedef reco::JetTagCollection JetTagCollection;
+
+ const edm::EDGetTokenT src_;
+ std::vector flav_names_;
+ std::vector input_names_;
+ std::vector output_names_;
+
+ bool skippedInference_ = false;
+};
+
+ParticleTransformerAK4SonicJetTagsProducer::ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet &iConfig)
+ : TritonEDProducer<>(iConfig),
+ src_(consumes(iConfig.getParameter("src"))),
+ flav_names_(iConfig.getParameter>("flav_names")),
+ input_names_(iConfig.getParameter>("input_names")),
+ output_names_(iConfig.getParameter>("output_names")) {
+ // get output names from flav_names
+ for (const auto &flav_name : flav_names_) {
+ produces(flav_name);
+ }
+}
+
+ParticleTransformerAK4SonicJetTagsProducer::~ParticleTransformerAK4SonicJetTagsProducer() {}
+
+void ParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
+ // pfParticleTransformerAK4JetTags
+ edm::ParameterSetDescription desc;
+ TritonClient::fillPSetDescription(desc);
+ desc.add("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
+ desc.add>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
+ desc.add>("output_names", {"softmax"});
+ desc.add>(
+ "flav_names", std::vector{"probb", "probbb", "problepb", "probc", "probuds", "probg"});
+
+ descriptions.add("pfParticleTransformerAK4SonicJetTags", desc);
+}
+
+void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEvent,
+ edm::EventSetup const &iSetup,
+ Input &iInput) {
+ edm::Handle tag_infos;
+ iEvent.getByToken(src_, tag_infos);
+ client_->setBatchSize(tag_infos->size());
+ skippedInference_ = false;
+ if (tag_infos->empty())
+ return;
+
+ // Find the max n_cpf, n_npf and n_vtx among all the jets in an event.
+ unsigned int max_n_cpf_counter = 0;
+ unsigned int max_n_npf_counter = 0;
+ unsigned int max_n_vtx_counter = 0;
+ for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
+ max_n_cpf_counter =
+ std::max(max_n_cpf_counter, static_cast(((*tag_infos)[jet_n]).features().c_pf_features.size()));
+ max_n_npf_counter =
+ std::max(max_n_npf_counter, static_cast(((*tag_infos)[jet_n]).features().n_pf_features.size()));
+ max_n_vtx_counter =
+ std::max(max_n_vtx_counter, static_cast(((*tag_infos)[jet_n]).features().sv_features.size()));
+ }
+
+ // If an event has no jet, or all jets has zero n_cpf, n_npf and n_vtx, the inference is skipped.
+ if (max_n_cpf_counter == 0 && max_n_npf_counter == 0 && max_n_vtx_counter == 0) {
+ client_->setBatchSize(0);
+ skippedInference_ = true;
+ return;
+ }
+
+ // all the jets in the same event will fill up the same amount of n_cpf, n_npf, n_vtx and send to server
+ const unsigned int target_n_cpf = std::clamp(max_n_cpf_counter, (unsigned int)1, (unsigned int)parT::n_cpf_accept);
+ const unsigned int target_n_npf = std::clamp(max_n_npf_counter, (unsigned int)1, (unsigned int)parT::n_npf_accept);
+ const unsigned int target_n_vtx = std::clamp(max_n_vtx_counter, (unsigned int)1, (unsigned int)parT::n_sv_accept);
+
+ const std::map target_n{{parT::kChargedCandidates, target_n_cpf},
+ {parT::kNeutralCandidates, target_n_npf},
+ {parT::kVertices, target_n_vtx},
+ {parT::kChargedCandidates4Vec, target_n_cpf},
+ {parT::kNeutralCandidates4Vec, target_n_npf},
+ {parT::kVertices4Vec, target_n_vtx}};
+
+ // loop through all groups of features
+ for (parT::InputFeatures ifeature = parT::kBegin; ifeature != parT::kEnd;
+ ifeature = static_cast(ifeature + 1)) {
+ const auto &group_name = input_names_[ifeature];
+ auto &input = iInput.at(group_name);
+
+ input.setShape(0, target_n.at(ifeature));
+ auto tdata = input.allocate(true);
+
+ for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
+ const auto &taginfo = (*tag_infos)[jet_n];
+ const auto &features = taginfo.features();
+ auto &vdata = (*tdata)[jet_n];
+ if (ifeature == parT::kChargedCandidates || ifeature == parT::kChargedCandidates4Vec)
+ parT_tensor_filler(vdata, ifeature, features.c_pf_features, target_n_cpf);
+ else if (ifeature == parT::kNeutralCandidates || ifeature == parT::kNeutralCandidates4Vec)
+ parT_tensor_filler(vdata, ifeature, features.n_pf_features, target_n_npf);
+ else if (ifeature == parT::kVertices || ifeature == parT::kVertices4Vec)
+ parT_tensor_filler(vdata, ifeature, features.sv_features, target_n_vtx);
+ }
+ input.toServer(tdata);
+ }
+}
+
+void ParticleTransformerAK4SonicJetTagsProducer::produce(edm::Event &iEvent,
+ const edm::EventSetup &iSetup,
+ Output const &iOutput) {
+ edm::Handle tag_infos;
+ iEvent.getByToken(src_, tag_infos);
+
+ // initialize output collection
+ std::vector> output_tags;
+ if (!tag_infos->empty()) {
+ auto jet_ref = tag_infos->begin()->jet();
+ auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
+ for (std::size_t i = 0; i < flav_names_.size(); i++) {
+ output_tags.emplace_back(std::make_unique(ref2prod));
+ }
+ } else {
+ for (std::size_t i = 0; i < flav_names_.size(); i++) {
+ output_tags.emplace_back(std::make_unique());
+ }
+ }
+ if (!tag_infos->empty()) {
+ if (!skippedInference_) {
+ const auto &output1 = iOutput.begin()->second;
+ const auto &outputs_from_server = output1.fromServer();
+
+ for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
+ const auto &taginfo = (*tag_infos)[jet_n];
+ const auto &jet_ref = tag_infos->at(jet_n).jet();
+
+ if (taginfo.features().is_filled) {
+ for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
+ (*(output_tags[flav_n]))[jet_ref] = outputs_from_server[jet_n][flav_n];
+ }
+ } else {
+ for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
+ (*(output_tags[flav_n]))[jet_ref] = -1.0;
+ }
+ }
+ }
+ } else {
+ for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
+ const auto &jet_ref = tag_infos->at(jet_n).jet();
+ for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
+ (*(output_tags[flav_n]))[jet_ref] = -1.0;
+ }
+ }
+ }
+ }
+ // put into the event
+ for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
+ iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
+ }
+}
+
+//define this as a plug-in
+DEFINE_FWK_MODULE(ParticleTransformerAK4SonicJetTagsProducer);
diff --git a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py
index 0fb1c8ceac050..4a7fd51782280 100644
--- a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py
+++ b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py
@@ -4,8 +4,26 @@
from RecoBTag.ONNXRuntime.pfParticleTransformerAK4JetTags_cfi import pfParticleTransformerAK4JetTags
from RecoBTag.ONNXRuntime.pfParticleTransformerAK4DiscriminatorsJetTags_cfi import pfParticleTransformerAK4DiscriminatorsJetTags
+from RecoBTag.ONNXRuntime.pfParticleTransformerAK4SonicJetTags_cfi import pfParticleTransformerAK4SonicJetTags as _pfParticleTransformerAK4SonicJetTags
from CommonTools.PileupAlgos.Puppi_cff import puppi
from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation
+from Configuration.ProcessModifiers.particleTransformerAK4SonicTriton_cff import particleTransformerAK4SonicTriton
+
+
+particleTransformerAK4SonicTriton.toReplaceWith(pfParticleTransformerAK4JetTags, _pfParticleTransformerAK4SonicJetTags.clone(
+ Client = cms.PSet(
+ timeout = cms.untracked.uint32(300),
+ mode = cms.string("Async"),
+ modelName = cms.string("particletransformer_AK4"),
+ modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particletransformer_AK4/config.pbtxt"),
+ modelVersion = cms.string(""),
+ verbose = cms.untracked.bool(False),
+ allowedTries = cms.untracked.uint32(0),
+ useSharedMemory = cms.untracked.bool(True),
+ compression = cms.untracked.string(""),
+ ),
+ flav_names = pfParticleTransformerAK4JetTags.flav_names,
+))
# declare all the discriminators
# probs
diff --git a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc
index da49b1b0119d7..9ecdf0302b45a 100644
--- a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc
+++ b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc
@@ -138,4 +138,81 @@ namespace btagbtvdeep {
*(++ptr) = neighbourTrack_features.dphi_PCAjetDirs;
}
+ std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features,
+ parT::InputFeatures ifeature) {
+ std::vector inputs;
+ if (ifeature == parT::kChargedCandidates) {
+ inputs.push_back(c_pf_features.btagPf_trackEtaRel);
+ inputs.push_back(c_pf_features.btagPf_trackPtRel);
+ inputs.push_back(c_pf_features.btagPf_trackPPar);
+ inputs.push_back(c_pf_features.btagPf_trackDeltaR);
+ inputs.push_back(c_pf_features.btagPf_trackPParRatio);
+ inputs.push_back(c_pf_features.btagPf_trackSip2dVal);
+ inputs.push_back(c_pf_features.btagPf_trackSip2dSig);
+ inputs.push_back(c_pf_features.btagPf_trackSip3dVal);
+ inputs.push_back(c_pf_features.btagPf_trackSip3dSig);
+ inputs.push_back(c_pf_features.btagPf_trackJetDistVal);
+ inputs.push_back(c_pf_features.ptrel);
+ inputs.push_back(c_pf_features.drminsv);
+ inputs.push_back(c_pf_features.vtx_ass);
+ inputs.push_back(c_pf_features.puppiw);
+ inputs.push_back(c_pf_features.chi2);
+ inputs.push_back(c_pf_features.quality);
+ } else if (ifeature == parT::kChargedCandidates4Vec) {
+ inputs.push_back(c_pf_features.px);
+ inputs.push_back(c_pf_features.py);
+ inputs.push_back(c_pf_features.pz);
+ inputs.push_back(c_pf_features.e);
+ }
+ return inputs;
+ }
+
+ std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features,
+ parT::InputFeatures ifeature) {
+ std::vector inputs;
+ if (ifeature == parT::kNeutralCandidates) {
+ inputs.push_back(n_pf_features.ptrel);
+ inputs.push_back(n_pf_features.etarel);
+ inputs.push_back(n_pf_features.phirel);
+ inputs.push_back(n_pf_features.deltaR);
+ inputs.push_back(n_pf_features.isGamma);
+ inputs.push_back(n_pf_features.hadFrac);
+ inputs.push_back(n_pf_features.drminsv);
+ inputs.push_back(n_pf_features.puppiw);
+ } else if (ifeature == parT::kNeutralCandidates4Vec) {
+ inputs.push_back(n_pf_features.px);
+ inputs.push_back(n_pf_features.py);
+ inputs.push_back(n_pf_features.pz);
+ inputs.push_back(n_pf_features.e);
+ }
+ return inputs;
+ }
+
+ std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features,
+ parT::InputFeatures ifeature) {
+ std::vector inputs;
+ if (ifeature == parT::kVertices) {
+ inputs.push_back(sv_features.pt);
+ inputs.push_back(sv_features.deltaR);
+ inputs.push_back(sv_features.mass);
+ inputs.push_back(sv_features.etarel);
+ inputs.push_back(sv_features.phirel);
+ inputs.push_back(sv_features.ntracks);
+ inputs.push_back(sv_features.chi2);
+ inputs.push_back(sv_features.normchi2);
+ inputs.push_back(sv_features.dxy);
+ inputs.push_back(sv_features.dxysig);
+ inputs.push_back(sv_features.d3d);
+ inputs.push_back(sv_features.d3dsig);
+ inputs.push_back(sv_features.costhetasvpv);
+ inputs.push_back(sv_features.enratio);
+ } else if (ifeature == parT::kVertices4Vec) {
+ inputs.push_back(sv_features.px);
+ inputs.push_back(sv_features.py);
+ inputs.push_back(sv_features.pz);
+ inputs.push_back(sv_features.e);
+ }
+ return inputs;
+ }
+
} // namespace btagbtvdeep