114):
115 '''TauJet identification inference based on ONNX models, for GNNs, transformers, etc...'''
116 acc = ComponentAccumulator()
117
118 try: id_flags = getattr(flags.Trigger.Offline.Tau, tau_id)
119 except NameError: raise ValueError(f'Invalid ONNX TauID configuration: {tau_id}')
120
121 if isinstance(id_flags.ONNXConfig, str):
122 network_config = {'NetworkFileInclusive': id_flags.ONNXConfig}
123 elif isinstance(id_flags.ONNXConfig, (list, tuple)) and len(id_flags.ONNXConfig) == 3:
124 network_config = {
125 'NetworkFile0P': id_flags.ONNXConfig[0],
126 'NetworkFile1P': id_flags.ONNXConfig[1],
127 'NetworkFile3P': id_flags.ONNXConfig[2],
128 }
129 else:
130 raise ValueError(f'Invalid {tau_id} ONNX network config file')
131
132 name = f'TrigTau_TauJetONNXEvaluator_{tau_id}'
133 if tau_container: name += f'_{tau_container}'
134 if hits_decoration_container: name += f'_{hits_decoration_container}'
135
136 ptau_decor = id_flags.OutputPTau if hasattr(id_flags, 'OutputPTau') else f'{tau_id}_ProbTau'
137 if ptau_sfx: ptau_decor += f'_{ptau_sfx}'
138 pjet_decor = id_flags.OutputPJet if hasattr(id_flags, 'OutputPJet') else f'{tau_id}_ProbJet'
139 if pjet_sfx: pjet_decor += f'_{pjet_sfx}'
140
141 acc.setPrivateTools(CompFactory.TauGNNEvaluator(
142 name = name,
143 useTRT = flags.Detector.EnableTRT,
144
145
146 **network_config,
147 InputLayerScalar = 'tau_vars',
148 InputLayerTracks = 'track_vars',
149 InputLayerClusters = 'cluster_vars',
150 InputLayerHits = 'hit_vars',
151 NodeNameTau = id_flags.NodeNameTau if hasattr(id_flags, 'NodeNameTau') else 'pTau',
152 NodeNameJet = id_flags.NodeNameJet if hasattr(id_flags, 'NodeNameJet') else 'pJet',
153
154
155 MaxTracks = id_flags.MaxTracks if hasattr(id_flags, 'MaxTracks') else 0,
156 MaxClusters = id_flags.MaxClusters if hasattr(id_flags, 'MaxClusters') else 0,
157 MaxClusterDR = 1.0,
158 MaxHits = id_flags.MaxHits if hasattr(id_flags, 'MaxHits') else 0,
159 TauContainerName = tau_container,
160 HitsHandleKey = hits_decoration_container,
161 VertexCorrection = False,
162 TrackClassification = False,
163
164
165 OutputVarname = f'{tau_id}_Score',
166 OutputDiscriminant = id_flags.OutputDiscriminant,
167 OutputPTau = ptau_decor,
168 OutputPJet = pjet_decor,
169
170
171 MinProngTrackPt = id_flags.MinProngTrackPt if hasattr(id_flags, 'MinProngTrackPt') else 0,
172 ))
173
174 return acc
175