From 9ccc8dacc5d926e302c38ab0bb70e803e08ecae0 Mon Sep 17 00:00:00 2001
From: Yao Yao <alexisyyao@gmail.com>
Date: Fri, 21 Jun 2024 15:10:58 -0400
Subject: [PATCH 1/2] Add ParT and UParT SONIC producer

---
 .../python/allSonicTriton_cff.py              |   5 +-
 .../particleTransformerAK4SonicTriton_cff.py  |   3 +
 ...edparticleTransformerAK4SonicTriton_cff.py |   3 +
 RecoBTag/ONNXRuntime/BuildFile.xml            |   1 +
 .../ONNXRuntime/interface/tensor_configs.h    |  83 +++++++
 .../ONNXRuntime/interface/tensor_fillers.h    |  77 ++++++
 ...rticleTransformerAK4ONNXJetTagsProducer.cc | 174 +++-----------
 ...ticleTransformerAK4SonicJetTagsProducer.cc | 187 +++++++++++++++
 ...rticleTransformerAK4ONNXJetTagsProducer.cc | 224 +++---------------
 ...ticleTransformerAK4SonicJetTagsProducer.cc | 200 ++++++++++++++++
 .../python/pfParticleTransformerAK4_cff.py    |  17 ++
 .../pfUnifiedParticleTransformerAK4_cff.py    |  18 ++
 RecoBTag/ONNXRuntime/src/tensor_fillers.cc    | 193 +++++++++++++++
 13 files changed, 854 insertions(+), 331 deletions(-)
 create mode 100644 Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py
 create mode 100644 Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py
 create mode 100644 RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc
 create mode 100644 RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc

diff --git a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py
index 396fcefd84a79..8f8df61d5be6d 100644
--- a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py
+++ b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py
@@ -5,6 +5,9 @@
 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
+from Configuration.ProcessModifiers.unifiedparticleTransformerAK4SonicTriton_cff import unifiedparticleTransformerAK4SonicTriton
+
 
 # collect all SonicTriton-related process modifiers here
-allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton)
+allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton,unifiedparticleTransformerAK4SonicTriton)
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/Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py b/Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py
new file mode 100644
index 0000000000000..d7b082a010613
--- /dev/null
+++ b/Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py
@@ -0,0 +1,3 @@
+import FWCore.ParameterSet.Config as cms
+
+unifiedparticleTransformerAK4SonicTriton = 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 @@
+<use name="PhysicsTools/ONNXRuntime"/>
 <use name="DataFormats/BTauReco"/>
 <export>
   <lib name="1"/>
diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h
index e528b4e92f74a..c5c14006cbbf2 100644
--- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h
+++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h
@@ -1,6 +1,8 @@
 #ifndef RecoBTag_ONNXRuntime_tensor_configs_h
 #define RecoBTag_ONNXRuntime_tensor_configs_h
 
+#include <map>
+
 namespace deepflavour {
 
   constexpr unsigned n_features_global = 15;
@@ -28,4 +30,85 @@ namespace deepvertex {
 
 }  // namespace deepvertex
 
+namespace parT {
+
+  enum InputFeatures {
+    kBegin=0,
+    kChargedCandidates=kBegin,
+    kNeutralCandidates=1,
+    kVertices=2,
+    kChargedCandidates4Vec=3,
+    kNeutralCandidates4Vec=4,
+    kVertices4Vec=5,
+    kEnd=6
+  };
+
+  constexpr unsigned n_cpf_accept = 25;
+  constexpr unsigned n_npf_accept = 25;
+  constexpr unsigned n_sv_accept = 5;
+
+  const std::map<InputFeatures, unsigned int> N_InputFeatures{
+    {kChargedCandidates, 16},
+    {kNeutralCandidates, 8},
+    {kVertices, 14},
+    {kChargedCandidates4Vec, 4},
+    {kNeutralCandidates4Vec, 4},
+    {kVertices4Vec, 4}
+  };
+
+  const std::map<InputFeatures, unsigned int> N_AcceptedFeatures{
+    {kChargedCandidates, n_cpf_accept},
+    {kNeutralCandidates, n_npf_accept},
+    {kVertices, n_sv_accept},
+    {kChargedCandidates4Vec, n_cpf_accept},
+    {kNeutralCandidates4Vec, n_npf_accept},
+    {kVertices4Vec, n_sv_accept}
+  };
+
+} // namespace parT
+
+
+namespace UparT {
+
+  enum InputFeatures {
+    kBegin=0,
+    kChargedCandidates=kBegin,
+    kLostTracks = 1,
+    kNeutralCandidates=2,
+    kVertices=3,
+    kChargedCandidates4Vec=4,
+    kLostTracks4Vec=5,
+    kNeutralCandidates4Vec=6,
+    kVertices4Vec=7,
+    kEnd=8
+  };
+
+  constexpr unsigned n_cpf_accept = 29;
+  constexpr unsigned n_lt_accept = 5;
+  constexpr unsigned n_npf_accept = 25;
+  constexpr unsigned n_sv_accept = 5;
+
+  const std::map<InputFeatures, unsigned int> N_InputFeatures{
+    {kChargedCandidates, 25},
+    {kLostTracks, 18},
+    {kNeutralCandidates, 8},
+    {kVertices, 14},
+    {kChargedCandidates4Vec, 4},
+    {kLostTracks4Vec, 4},
+    {kNeutralCandidates4Vec, 4},
+    {kVertices4Vec, 4}
+  };
+
+  const std::map<InputFeatures, unsigned int> N_AcceptedFeatures{
+    {kChargedCandidates, n_cpf_accept},
+    {kLostTracks, n_lt_accept},
+    {kNeutralCandidates, n_npf_accept},
+    {kVertices, n_sv_accept},
+    {kChargedCandidates4Vec, n_cpf_accept},
+    {kLostTracks4Vec, n_lt_accept},
+    {kNeutralCandidates4Vec, n_npf_accept},
+    {kVertices4Vec, n_sv_accept}
+  };
+
+} // namespace UparT
 #endif
diff --git a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h
index bb1a25a2d0f0b..62100573bcd5b 100644
--- a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h
+++ b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h
@@ -2,6 +2,10 @@
 #define RecoBTag_ONNXRuntime_tensor_fillers_h
 
 #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h"
+#include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
+#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h"
+#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
+#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
 
 namespace btagbtvdeep {
 
@@ -19,6 +23,79 @@ namespace btagbtvdeep {
 
   void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features);
 
+  std::vector<float> inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature);
+
+  std::vector<float> inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature);
+
+  std::vector<float> inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, UparT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::LostTracksFeatures& lt_features, UparT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, UparT::InputFeatures ifeature);
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, UparT::InputFeatures ifeature);
+
+  template<class parT_features>
+  void parT_tensor_filler(cms::Ort::FloatArrays& data, const parT::InputFeatures ifeature, const std::vector<parT_features>& 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<float>& 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<class parT_features>
+  void parT_tensor_filler(std::vector<float>& vdata, const parT::InputFeatures ifeature, const std::vector<parT_features>& 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<float>& 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
+  }
+
+  template<class UparT_features>
+  void UparT_tensor_filler(cms::Ort::FloatArrays& data, const UparT::InputFeatures ifeature, const std::vector<UparT_features>& 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 * UparT::N_InputFeatures.at(ifeature)];
+      start = ptr;
+      const std::vector<float>& inputs = inputs_UparT(f, ifeature);
+      for (unsigned int i = 0; i < inputs.size(); i++) {
+        *ptr = inputs[i];
+        ++ptr;
+      }
+      if (inputs.size() > 0) --ptr;
+      assert(start + UparT::N_InputFeatures.at(ifeature) - 1 == ptr);
+    }
+  }
+
+  template<class UparT_features>
+  void UparT_tensor_filler(std::vector<float>& vdata, const UparT::InputFeatures ifeature, const std::vector<UparT_features>& features, const unsigned int target_n) {
+    unsigned int n = std::clamp((unsigned int)features.size(), (unsigned int)0, (unsigned int)UparT::N_AcceptedFeatures.at(ifeature));
+    for (unsigned int count = 0; count < n; count++) {
+      const std::vector<float>& inputs = inputs_UparT(features.at(count), ifeature);
+      vdata.insert(vdata.end(), inputs.begin(), inputs.end());
+    }
+    unsigned int n_features = UparT::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..e63b5504936b4 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<edm::GlobalCache<ONNXRuntime>> {
@@ -27,7 +30,7 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer
 
   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet&);
   static void globalEndJob(const ONNXRuntime*);
-
+ 
 private:
   typedef std::vector<reco::ParticleTransformerAK4TagInfo> TagInfoCollection;
   typedef reco::JetTagCollection JetTagCollection;
@@ -41,24 +44,9 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer
   std::vector<std::string> flav_names_;
   std::vector<std::string> input_names_;
   std::vector<std::string> 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<unsigned> input_sizes_;
   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
 
@@ -84,7 +72,7 @@ void ParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Configurat
   desc.add<edm::InputTag>("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
   desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
   desc.add<edm::FileInPath>("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<std::vector<std::string>>("output_names", {"softmax"});
   desc.add<std::vector<std::string>>(
       "flav_names", std::vector<std::string>{"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<btagbtvdeep::ParticleTransformerAK4Features> 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::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::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::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,25 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
 }
 
 void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::ParticleTransformerAK4Features features) {
-  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..cf515d170373e
--- /dev/null
+++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc
@@ -0,0 +1,187 @@
+#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<reco::ParticleTransformerAK4TagInfo> TagInfoCollection;
+  typedef reco::JetTagCollection JetTagCollection;
+
+  const edm::EDGetTokenT<TagInfoCollection> src_;
+  std::vector<std::string> flav_names_;
+  std::vector<std::string> input_names_;
+  std::vector<std::string> output_names_;
+
+  bool skippedInference_ = false;
+};
+
+ParticleTransformerAK4SonicJetTagsProducer::ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet& iConfig)
+    : TritonEDProducer<>(iConfig),
+      src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
+      flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
+      input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
+      output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
+  // get output names from flav_names
+  for (const auto& flav_name : flav_names_) {
+    produces<JetTagCollection>(flav_name);
+  }
+}
+
+ParticleTransformerAK4SonicJetTagsProducer::~ParticleTransformerAK4SonicJetTagsProducer() {}
+
+void ParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  // pfParticleTransformerAK4JetTags
+  edm::ParameterSetDescription desc;
+  TritonClient::fillPSetDescription(desc);
+  desc.add<edm::InputTag>("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
+  desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
+  desc.add<std::vector<std::string>>("output_names", {"softmax"});
+  desc.add<std::vector<std::string>>(
+      "flav_names", std::vector<std::string>{"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<TagInfoCollection> 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<unsigned int>(((*tag_infos)[jet_n]).features().c_pf_features.size()));
+    max_n_npf_counter = std::max(max_n_npf_counter,
+                         static_cast<unsigned int>(((*tag_infos)[jet_n]).features().n_pf_features.size()));
+    max_n_vtx_counter = std::max(max_n_vtx_counter,
+                         static_cast<unsigned int>(((*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<parT::InputFeatures, unsigned int> 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<parT::InputFeatures>(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<float>(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<TagInfoCollection> tag_infos;
+  iEvent.getByToken(src_, tag_infos);
+
+  // initialize output collection
+  std::vector<std::unique_ptr<JetTagCollection>> 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<JetTagCollection>(ref2prod));
+    }
+  } else {
+    for (std::size_t i = 0; i < flav_names_.size(); i++) {
+      output_tags.emplace_back(std::make_unique<JetTagCollection>());
+    }
+  }
+  if (!tag_infos->empty()) {
+    if (!skippedInference_) {
+      const auto &output1 = iOutput.begin()->second;
+      const auto &outputs_from_server = output1.fromServer<float>();
+
+      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/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc
index a6864da0c3de5..2f9548d90f880 100644
--- a/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc
+++ b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.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 UnifiedParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
@@ -42,28 +45,10 @@ class UnifiedParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDP
   std::vector<std::string> input_names_;
   std::vector<std::string> output_names_;
 
-  enum InputIndexes {
-    kChargedCandidates = 0,
-    kLostTracks = 1,
-    kNeutralCandidates = 2,
-    kVertices = 3,
-    kChargedCandidates4Vec = 4,
-    kLostTracks4Vec = 5,
-    kNeutralCandidates4Vec = 6,
-    kVertices4Vec = 7
-  };
   unsigned n_cpf_;
-  constexpr static unsigned n_features_cpf_ = 25;
-  constexpr static unsigned n_pairwise_features_cpf_ = 4;
   unsigned n_lt_;
-  constexpr static unsigned n_features_lt_ = 18;
-  constexpr static unsigned n_pairwise_features_lt_ = 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;
   std::vector<unsigned> input_sizes_;
   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
 
@@ -89,7 +74,7 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Con
   desc.add<edm::InputTag>("src", edm::InputTag("pfUnifiedParticleTransformerAK4TagInfos"));
   desc.add<std::vector<std::string>>(
       "input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6", "input_7", "input_8"});
-  desc.add<edm::FileInPath>("model_path", edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V00/UParTAK4.onnx"));
+  desc.add<edm::FileInPath>("model_path", edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V00/modelfile/model.onnx"));
   desc.add<std::vector<std::string>>("output_names", {"softmax"});
   desc.add<std::vector<std::string>>(
       "flav_names",
@@ -135,14 +120,14 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEven
       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_lt_, (int64_t)n_features_lt_},
-                       {(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_lt_, (int64_t)n_pairwise_features_lt_},
-                       {(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)UparT::N_InputFeatures.at(UparT::kChargedCandidates)},
+                       {(int64_t)1, (int64_t)n_lt_, (int64_t)UparT::N_InputFeatures.at(UparT::kLostTracks)},
+                       {(int64_t)1, (int64_t)n_npf_, (int64_t)UparT::N_InputFeatures.at(UparT::kNeutralCandidates)},
+                       {(int64_t)1, (int64_t)n_sv_, (int64_t)UparT::N_InputFeatures.at(UparT::kVertices)},
+                       {(int64_t)1, (int64_t)n_cpf_, (int64_t)UparT::N_InputFeatures.at(UparT::kChargedCandidates4Vec)},
+                       {(int64_t)1, (int64_t)n_lt_, (int64_t)UparT::N_InputFeatures.at(UparT::kLostTracks4Vec)},
+                       {(int64_t)1, (int64_t)n_npf_, (int64_t)UparT::N_InputFeatures.at(UparT::kNeutralCandidates4Vec)},
+                       {(int64_t)1, (int64_t)n_sv_, (int64_t)UparT::N_InputFeatures.at(UparT::kVertices4Vec)}};
 
       outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0];
       assert(outputs.size() == flav_names_.size());
@@ -165,20 +150,20 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
   const auto& features = taginfo.features();
 
   /// We require a fixed size due to an ONNX conversion issue (to be improved in the future ?) ///
-  n_cpf_ = (unsigned int)29;
-  n_lt_ = (unsigned int)5;
-  n_npf_ = (unsigned int)25;
-  n_sv_ = (unsigned int)5;
+  n_cpf_ = (unsigned int)UparT::n_cpf_accept;
+  n_lt_ = (unsigned int)UparT::n_lt_accept;
+  n_npf_ = (unsigned int)UparT::n_npf_accept;
+  n_sv_ = (unsigned int)UparT::n_sv_accept;
 
   input_sizes_ = {
-      n_cpf_ * n_features_cpf_,
-      n_lt_ * n_features_lt_,
-      n_npf_ * n_features_npf_,
-      n_sv_ * n_features_sv_,
-      n_cpf_ * n_pairwise_features_cpf_,
-      n_lt_ * n_pairwise_features_lt_,
-      n_npf_ * n_pairwise_features_npf_,
-      n_sv_ * n_pairwise_features_sv_,
+      n_cpf_ * UparT::N_InputFeatures.at(UparT::kChargedCandidates),
+      n_lt_ * UparT::N_InputFeatures.at(UparT::kLostTracks),
+      n_npf_ * UparT::N_InputFeatures.at(UparT::kNeutralCandidates),
+      n_sv_ * UparT::N_InputFeatures.at(UparT::kVertices),
+      n_cpf_ * UparT::N_InputFeatures.at(UparT::kChargedCandidates4Vec),
+      n_lt_ * UparT::N_InputFeatures.at(UparT::kLostTracks4Vec),
+      n_npf_ * UparT::N_InputFeatures.at(UparT::kNeutralCandidates4Vec),
+      n_sv_ * UparT::N_InputFeatures.at(UparT::kVertices4Vec),
   };
   // init data storage
   data_.clear();
@@ -191,167 +176,30 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
 
 void UnifiedParticleTransformerAK4ONNXJetTagsProducer::make_inputs(
     btagbtvdeep::UnifiedParticleTransformerAK4Features features) {
-  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;
-    *(++ptr) = c_pf_features.charge;
-    *(++ptr) = c_pf_features.dz;
-    *(++ptr) = c_pf_features.btagPf_trackDecayLen;
-    *(++ptr) = c_pf_features.HadFrac;
-    *(++ptr) = c_pf_features.CaloFrac;
-    *(++ptr) = c_pf_features.pdgID;
-    *(++ptr) = c_pf_features.lostInnerHits;
-    *(++ptr) = c_pf_features.numberOfPixelHits;
-    *(++ptr) = c_pf_features.numberOfStripHits;
-
-    assert(start + n_features_cpf_ - 1 == ptr);
-  }
-
-  // n_lt candidates
   auto max_lt_n = std::min(features.lt_features.size(), (std::size_t)n_lt_);
-  for (std::size_t lt_n = 0; lt_n < max_lt_n; lt_n++) {
-    const auto& lt_features = features.lt_features.at(lt_n);
-    ptr = &data_[kLostTracks][offset + lt_n * n_features_lt_];
-    start = ptr;
-    *ptr = lt_features.btagPf_trackEtaRel;
-    *(++ptr) = lt_features.btagPf_trackPtRel;
-    *(++ptr) = lt_features.btagPf_trackPPar;
-    *(++ptr) = lt_features.btagPf_trackDeltaR;
-    *(++ptr) = lt_features.btagPf_trackPParRatio;
-    *(++ptr) = lt_features.btagPf_trackSip2dVal;
-    *(++ptr) = lt_features.btagPf_trackSip2dSig;
-    *(++ptr) = lt_features.btagPf_trackSip3dVal;
-    *(++ptr) = lt_features.btagPf_trackSip3dSig;
-    *(++ptr) = lt_features.btagPf_trackJetDistVal;
-    *(++ptr) = lt_features.drminsv;
-    *(++ptr) = lt_features.charge;
-    *(++ptr) = lt_features.puppiw;
-    *(++ptr) = lt_features.chi2;
-    *(++ptr) = lt_features.quality;
-    *(++ptr) = lt_features.lostInnerHits;
-    *(++ptr) = lt_features.numberOfPixelHits;
-    *(++ptr) = lt_features.numberOfStripHits;
-    assert(start + n_features_lt_ - 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
+  UparT_tensor_filler(data_, UparT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset);
+  // lt candidates
+  UparT_tensor_filler(data_, UparT::kLostTracks, features.lt_features, max_lt_n, start, offset);
+  // n_pf candidates
+  UparT_tensor_filler(data_, UparT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset);
+  // sv candidates
+  UparT_tensor_filler(data_, UparT::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);
-  }
-
+  UparT_tensor_filler(data_, UparT::kChargedCandidates4Vec, features.c_pf_features, max_c_pf_n, start, offset);
   // lt pairwise features (4-vectors) specific case requiring (pt,eta,phi,e)
-  auto max_lt_N = std::min(features.lt_features.size(), (std::size_t)n_lt_);
-  for (std::size_t lt_N = 0; lt_N < max_lt_N; lt_N++) {
-    const auto& lt_pairwise_features = features.lt_features.at(lt_N);
-    ptr = &data_[kLostTracks4Vec][offset + lt_N * n_pairwise_features_lt_];
-    start = ptr;
-    *ptr = lt_pairwise_features.pt;
-    *(++ptr) = lt_pairwise_features.eta;
-    *(++ptr) = lt_pairwise_features.phi;
-    *(++ptr) = lt_pairwise_features.e;
-
-    assert(start + n_pairwise_features_lt_ - 1 == ptr);
-  }
-
+  UparT_tensor_filler(data_, UparT::kLostTracks4Vec, features.lt_features, max_lt_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);
-  }
-
+  UparT_tensor_filler(data_, UparT::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);
-  }
+  UparT_tensor_filler(data_, UparT::kVertices4Vec, features.sv_features, max_sv_n, start, offset);
 }
 
 //define this as a plug-in
diff --git a/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc
new file mode 100644
index 0000000000000..49d34a9142a27
--- /dev/null
+++ b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc
@@ -0,0 +1,200 @@
+#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/UnifiedParticleTransformerAK4TagInfo.h"
+#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.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 UnifiedParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> {
+public:
+  explicit UnifiedParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet&);
+  ~UnifiedParticleTransformerAK4SonicJetTagsProducer() 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<reco::UnifiedParticleTransformerAK4TagInfo> TagInfoCollection;
+  typedef reco::JetTagCollection JetTagCollection;
+
+  const edm::EDGetTokenT<TagInfoCollection> src_;
+  std::vector<std::string> flav_names_;
+  std::vector<std::string> input_names_;
+  std::vector<std::string> output_names_;
+
+  bool skippedInference_ = false;
+};
+
+UnifiedParticleTransformerAK4SonicJetTagsProducer::UnifiedParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet& iConfig)
+    : TritonEDProducer<>(iConfig),
+      src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
+      flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
+      input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
+      output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
+  // get output names from flav_names
+  for (const auto& flav_name : flav_names_) {
+    produces<JetTagCollection>(flav_name);
+  }
+}
+
+UnifiedParticleTransformerAK4SonicJetTagsProducer::~UnifiedParticleTransformerAK4SonicJetTagsProducer() {}
+
+void UnifiedParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  // pfUnifiedParticleTransformerAK4JetTags
+  edm::ParameterSetDescription desc;
+  TritonClient::fillPSetDescription(desc);
+  desc.add<edm::InputTag>("src", edm::InputTag("pfUnifiedParticleTransformerAK4TagInfos"));
+  desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6", "input_7", "input_8"});
+  desc.add<std::vector<std::string>>("output_names", {"softmax"});
+  desc.add<std::vector<std::string>>(
+      "flav_names", 
+      std::vector<std::string>{"probb",        "probbb",       "problepb",     "probc",         "probs",
+                               "probu",        "probd",        "probg",        "probele",       "probmu",
+                               "probtaup1h0p", "probtaup1h1p", "probtaup1h2p", "probtaup3h0p",  "probtaup3h1p",
+                               "probtaum1h0p", "probtaum1h1p", "probtaum1h2p", "probtaum3h0p",  "probtaum3h1p",
+                               "ptcorr",       "ptreshigh",    "ptreslow",     "ptnu",          "probemudata",
+                               "probemumc",    "probdimudata", "probdimumc",   "probmutaudata", "probmutaumc"});
+
+  descriptions.add("pfUnifiedParticleTransformerAK4SonicJetTags", desc);
+}
+
+void UnifiedParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) {
+  edm::Handle<TagInfoCollection> 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_lt_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<unsigned int>(((*tag_infos)[jet_n]).features().c_pf_features.size()));
+    max_n_lt_counter = std::max(max_n_lt_counter, 
+                         static_cast<unsigned int>(((*tag_infos)[jet_n]).features().lt_features.size()));
+    max_n_npf_counter = std::max(max_n_npf_counter,
+                         static_cast<unsigned int>(((*tag_infos)[jet_n]).features().n_pf_features.size()));
+    max_n_vtx_counter = std::max(max_n_vtx_counter,
+                         static_cast<unsigned int>(((*tag_infos)[jet_n]).features().sv_features.size()));
+  }
+  
+  // If an event has no jet, or all jets has zero n_cpf, n_lt, n_npf and n_vtx, the inference is skipped.
+  if (max_n_cpf_counter == 0 && max_n_lt_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 = (unsigned int)UparT::n_cpf_accept;
+  const unsigned int target_n_lt = (unsigned int)UparT::n_lt_accept;
+  const unsigned int target_n_npf = (unsigned int)UparT::n_npf_accept;
+  const unsigned int target_n_vtx = (unsigned int)UparT::n_sv_accept;
+
+  const std::map<UparT::InputFeatures, unsigned int> target_n {
+    {UparT::kChargedCandidates,     target_n_cpf},
+    {UparT::kLostTracks,            target_n_lt},
+    {UparT::kNeutralCandidates,     target_n_npf},
+    {UparT::kVertices,              target_n_vtx},
+    {UparT::kChargedCandidates4Vec, target_n_cpf},
+    {UparT::kLostTracks4Vec,        target_n_lt},
+    {UparT::kNeutralCandidates4Vec, target_n_npf},
+    {UparT::kVertices4Vec,          target_n_vtx}
+  };
+ 
+  // loop through all groups of features
+  for (UparT::InputFeatures ifeature = UparT::kBegin; ifeature != UparT::kEnd; ifeature = static_cast<UparT::InputFeatures>(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<float>(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 == UparT::kChargedCandidates || ifeature == UparT::kChargedCandidates4Vec) 
+        UparT_tensor_filler(vdata, ifeature, features.c_pf_features, target_n_cpf);
+      else if (ifeature == UparT::kNeutralCandidates || ifeature == UparT::kNeutralCandidates4Vec) 
+        UparT_tensor_filler(vdata, ifeature, features.n_pf_features, target_n_npf);
+      else if (ifeature == UparT::kVertices || ifeature == UparT::kVertices4Vec)
+        UparT_tensor_filler(vdata, ifeature, features.sv_features, target_n_vtx);
+    }
+    input.toServer(tdata);
+  }
+}
+
+void UnifiedParticleTransformerAK4SonicJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const &iOutput) {
+  edm::Handle<TagInfoCollection> tag_infos;
+  iEvent.getByToken(src_, tag_infos);
+
+  // initialize output collection
+  std::vector<std::unique_ptr<JetTagCollection>> 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<JetTagCollection>(ref2prod));
+    }
+  } else {
+    for (std::size_t i = 0; i < flav_names_.size(); i++) {
+      output_tags.emplace_back(std::make_unique<JetTagCollection>());
+    }
+  }
+  if (!tag_infos->empty()) {
+    if (!skippedInference_) {
+      const auto &output1 = iOutput.begin()->second;
+      const auto &outputs_from_server = output1.fromServer<float>();
+
+      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(UnifiedParticleTransformerAK4SonicJetTagsProducer);
diff --git a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py
index 9df64c14f3406..b721636164802 100644
--- a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py
+++ b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py
@@ -6,6 +6,23 @@
 from RecoBTag.ONNXRuntime.pfParticleTransformerAK4DiscriminatorsJetTags_cfi import pfParticleTransformerAK4DiscriminatorsJetTags
 from CommonTools.PileupAlgos.Puppi_cff import puppi
 from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation
+from RecoBTag.ONNXRuntime.pfParticleTransformerAK4SonicJetTags_cfi import pfParticleTransformerAK4SonicJetTags as _pfParticleTransformerAK4SonicJetTags
+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"), # double check
+        modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particletransformer_AK4/config.pbtxt"), # this is SONIC, not currently in the CMSSW, so the models/ will be copied to this location privately
+        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/python/pfUnifiedParticleTransformerAK4_cff.py b/RecoBTag/ONNXRuntime/python/pfUnifiedParticleTransformerAK4_cff.py
index aee089f50ef68..b9584d00401bc 100644
--- a/RecoBTag/ONNXRuntime/python/pfUnifiedParticleTransformerAK4_cff.py
+++ b/RecoBTag/ONNXRuntime/python/pfUnifiedParticleTransformerAK4_cff.py
@@ -6,6 +6,24 @@
 from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4DiscriminatorsJetTags_cfi import pfUnifiedParticleTransformerAK4DiscriminatorsJetTags
 from CommonTools.PileupAlgos.Puppi_cff import puppi
 from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation
+from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4SonicJetTags_cfi import pfUnifiedParticleTransformerAK4SonicJetTags as _pfUnifiedParticleTransformerAK4SonicJetTags
+from Configuration.ProcessModifiers.unifiedparticleTransformerAK4SonicTriton_cff import unifiedparticleTransformerAK4SonicTriton
+
+unifiedparticleTransformerAK4SonicTriton.toReplaceWith(pfUnifiedParticleTransformerAK4JetTags, _pfUnifiedParticleTransformerAK4SonicJetTags.clone(
+    Client = cms.PSet(
+        timeout = cms.untracked.uint32(300),
+        mode = cms.string("Async"),
+        modelName = cms.string("unifiedparticletransformer_AK4"), # double check
+        modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/unifiedparticletransformer_AK4/config.pbtxt"), # this is SONIC, not currently in the CMSSW, so the models/ will be copied to this location privately
+        modelVersion = cms.string(""),
+        verbose = cms.untracked.bool(False),
+        allowedTries = cms.untracked.uint32(0),
+        useSharedMemory = cms.untracked.bool(True),
+        compression = cms.untracked.string(""),
+    ),
+    flav_names = pfUnifiedParticleTransformerAK4JetTags.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..3a5f0357ea6fc 100644
--- a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc
+++ b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc
@@ -138,4 +138,197 @@ namespace btagbtvdeep {
     *(++ptr) = neighbourTrack_features.dphi_PCAjetDirs;
   }
 
+  std::vector<float> inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature) {
+    std::vector<float> 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<float> inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature) {
+    std::vector<float> 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<float> inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature) {
+    std::vector<float> 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;
+  }
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, UparT::InputFeatures ifeature) {
+    std::vector<float> inputs;
+    if (ifeature == UparT::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);
+      inputs.push_back(c_pf_features.charge);
+      inputs.push_back(c_pf_features.dz);
+      inputs.push_back(c_pf_features.btagPf_trackDecayLen);
+      inputs.push_back(c_pf_features.HadFrac);
+      inputs.push_back(c_pf_features.CaloFrac);
+      inputs.push_back(c_pf_features.pdgID);
+      inputs.push_back(c_pf_features.lostInnerHits);
+      inputs.push_back(c_pf_features.numberOfPixelHits);
+      inputs.push_back(c_pf_features.numberOfStripHits);
+    }
+    else if (ifeature == UparT::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<float> inputs_UparT(const btagbtvdeep::LostTracksFeatures& lt_features, UparT::InputFeatures ifeature) {
+    std::vector<float> inputs;
+    if (ifeature == UparT::kLostTracks) {
+      inputs.push_back(lt_features.btagPf_trackEtaRel);
+      inputs.push_back(lt_features.btagPf_trackPtRel);
+      inputs.push_back(lt_features.btagPf_trackPPar);
+      inputs.push_back(lt_features.btagPf_trackDeltaR);
+      inputs.push_back(lt_features.btagPf_trackPParRatio);
+      inputs.push_back(lt_features.btagPf_trackSip2dVal);
+      inputs.push_back(lt_features.btagPf_trackSip2dSig);
+      inputs.push_back(lt_features.btagPf_trackSip3dVal);
+      inputs.push_back(lt_features.btagPf_trackSip3dSig);
+      inputs.push_back(lt_features.btagPf_trackJetDistVal);
+      inputs.push_back(lt_features.drminsv);
+      inputs.push_back(lt_features.charge);
+      inputs.push_back(lt_features.puppiw);
+      inputs.push_back(lt_features.chi2);
+      inputs.push_back(lt_features.quality);
+      inputs.push_back(lt_features.lostInnerHits);
+      inputs.push_back(lt_features.numberOfPixelHits);
+      inputs.push_back(lt_features.numberOfStripHits);
+    }   
+    else if (ifeature == UparT::kLostTracks4Vec) {
+      inputs.push_back(lt_features.pt);
+      inputs.push_back(lt_features.eta);
+      inputs.push_back(lt_features.phi);
+      inputs.push_back(lt_features.e);
+    }   
+    return inputs;
+  }
+
+  std::vector<float> inputs_UparT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, UparT::InputFeatures ifeature) {
+    std::vector<float> inputs;
+    if (ifeature == UparT::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 == UparT::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<float> inputs_UparT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, UparT::InputFeatures ifeature) {
+    std::vector<float> inputs;
+    if (ifeature == UparT::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 == UparT::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

From 75250f9424ef383b60262351e58075690fb407ae Mon Sep 17 00:00:00 2001
From: Yao Yao <alexisyyao@gmail.com>
Date: Fri, 26 Jul 2024 11:03:18 -0400
Subject: [PATCH 2/2] fix a bug in producer

---
 .../UnifiedParticleTransformerAK4SonicJetTagsProducer.cc        | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc
index 49d34a9142a27..bc7810f159dd4 100644
--- a/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc
+++ b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc
@@ -136,6 +136,8 @@ void UnifiedParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const
       auto &vdata = (*tdata)[jet_n];
       if (ifeature == UparT::kChargedCandidates || ifeature == UparT::kChargedCandidates4Vec) 
         UparT_tensor_filler(vdata, ifeature, features.c_pf_features, target_n_cpf);
+      else if (ifeature == UparT::kLostTracks || ifeature == UparT::kLostTracks4Vec)
+        UparT_tensor_filler(vdata, ifeature, features.lt_features, target_n_lt);
       else if (ifeature == UparT::kNeutralCandidates || ifeature == UparT::kNeutralCandidates4Vec) 
         UparT_tensor_filler(vdata, ifeature, features.n_pf_features, target_n_npf);
       else if (ifeature == UparT::kVertices || ifeature == UparT::kVertices4Vec)