Jet Flavour Extension

.gitignore

@@ -36,3 +36,4 @@ benchmark/*
 /profile/*
 /statprof/*
 /debug/*
+/examples/flavour-tagging/data

Project.toml

@@ -17,10 +17,13 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 [weakdeps]
 EDM4hep = "eb32b910-dde9-4347-8fce-cd6be3498f0c"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+ONNXRunTime = "e034b28e-924e-41b2-b98f-d2bbeb830c6a"
+JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 
 [extensions]
 EDM4hepJets = "EDM4hep"
 JetVisualisation = "Makie"
+JetFlavourTagging = ["EDM4hep", "ONNXRunTime", "JSON"]
 
 [compat]
 Accessors = "0.1.36"
@@ -35,6 +38,7 @@ LoopVectorization = "0.12.170"
 LorentzVectorHEP = "0.1.6"
 Makie = "0.20, 0.21, 0.22"
 MuladdMacro = "0.2.4"
+ONNXRunTime = "1.3.1"
 StructArrays = "0.6.18, 0.7"
 Test = "1.9"
 julia = "1.9"

examples/flavour-tagging/Project.toml

@@ -0,0 +1,20 @@
+[deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
+CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2"
+EDM4hep = "eb32b910-dde9-4347-8fce-cd6be3498f0c"
+EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
+IJulia = "7073ff75-c697-5162-941a-fcdaad2a7d2a"
+JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+JetReconstruction = "44e8cb2c-dfab-4825-9c70-d4808a591196"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
+LorentzVectorHEP = "f612022c-142a-473f-8cfd-a09cf3793c6c"
+LorentzVectors = "3f54b04b-17fc-5cd4-9758-90c048d965e3"
+MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
+ONNXRunTime = "e034b28e-924e-41b2-b98f-d2bbeb830c6a"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+
+[sources]
+JetReconstruction = {path = "../.."}

examples/flavour-tagging/README.md

@@ -0,0 +1,10 @@
+# Jet Reconstruction Flavour Tagging Module Examples
+
+## `simple-flavour-tagging.jl`
+
+This script will perform a simple flavour tagging with one events only.
+It will use the `FlavourTagging` module to tag the jets in the event and print the results.
+
+## `simple-flavour-tagging.ipynb`
+
+The same as the above example, but using a Jupyter notebook.

examples/flavour-tagging/simple-flavour-tagging.ipynb

@@ -0,0 +1,293 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "21a84728",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "using EDM4hep\n",
+    "using EDM4hep.RootIO\n",
+    "using StaticArrays\n",
+    "using LorentzVectorHEP\n",
+    "using JSON\n",
+    "using ONNXRunTime\n",
+    "using StructArrays"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "3b92ca23",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Project.toml`\n",
+      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Manifest.toml`\n",
+      "\u001b[32m\u001b[1m   Resolving\u001b[22m\u001b[39m package versions...\n",
+      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Project.toml`\n",
+      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.11/Manifest.toml`\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m\u001b[1mStatus\u001b[22m\u001b[39m `~/.julia/environments/v1.11/Project.toml`\n",
+      "  \u001b[90m[c7e460c6] \u001b[39mArgParse v1.2.0\n",
+      "  \u001b[90m[a93c6f00] \u001b[39mDataFrames v1.7.0\n",
+      "  \u001b[90m[eb32b910] \u001b[39mEDM4hep v0.4.3\n",
+      "\u001b[32m⌃\u001b[39m \u001b[90m[5c1252a2] \u001b[39mGeometryBasics v0.5.7\n",
+      "  \u001b[90m[7073ff75] \u001b[39mIJulia v1.27.0\n",
+      "  \u001b[90m[682c06a0] \u001b[39mJSON v0.21.4\n",
+      "  \u001b[90m[44e8cb2c] \u001b[39mJetReconstruction v0.5.0-dev `~/Developer/2025/JuliaHEPForkToMain/JetReconstruction.jl`\n",
+      "  \u001b[90m[98e50ef6] \u001b[39mJuliaFormatter v2.1.2\n",
+      "\u001b[32m⌃\u001b[39m \u001b[90m[f612022c] \u001b[39mLorentzVectorHEP v0.1.6\n",
+      "  \u001b[90m[e034b28e] \u001b[39mONNXRunTime v1.3.1\n",
+      "  \u001b[90m[91a5bcdd] \u001b[39mPlots v1.40.13\n",
+      "  \u001b[90m[90137ffa] \u001b[39mStaticArrays v1.9.13\n",
+      "  \u001b[90m[10745b16] \u001b[39mStatistics v1.11.1\n",
+      "  \u001b[90m[2913bbd2] \u001b[39mStatsBase v0.34.5\n",
+      "\u001b[33m⌅\u001b[39m \u001b[90m[09ab397b] \u001b[39mStructArrays v0.6.21\n",
+      "\u001b[32m⌃\u001b[39m \u001b[90m[1e6cf692] \u001b[39mTestEnv v1.102.0\n",
+      "  \u001b[90m[3cd96dde] \u001b[39mUnROOT v0.10.36\n",
+      "  \u001b[90m[37e2e46d] \u001b[39mLinearAlgebra v1.11.0\n",
+      "\u001b[36m\u001b[1mInfo\u001b[22m\u001b[39m Packages marked with \u001b[32m⌃\u001b[39m and \u001b[33m⌅\u001b[39m have new versions available. Those with \u001b[32m⌃\u001b[39m may be upgradable, but those with \u001b[33m⌅\u001b[39m are restricted by compatibility constraints from upgrading. To see why use `status --outdated`\n"
+     ]
+    }
+   ],
+   "source": [
+    "using Pkg\n",
+    "Pkg.resolve()\n",
+    "Pkg.develop(Pkg.PackageSpec(path = \"/Users/harrywanghc/Developer/2025/JuliaHEPForkToMain/JetReconstruction.jl\"))\n",
+    "using JetReconstruction\n",
+    "Pkg.status()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "379f4304",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The model predicts these flavor classes:\n",
+      " - recojet_isG\n",
+      " - recojet_isQ\n",
+      " - recojet_isS\n",
+      " - recojet_isC\n",
+      " - recojet_isB\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Paths to model files\n",
+    "model_dir = \"data/wc_pt_7classes_12_04_2023\"\n",
+    "onnx_path = joinpath(model_dir, \"fccee_flavtagging_edm4hep_wc_v1.onnx\")\n",
+    "json_path = joinpath(model_dir, \"fccee_flavtagging_edm4hep_wc_v1.json\")\n",
+    "\n",
+    "# Load the configuration and model\n",
+    "config = JSON.parsefile(json_path)\n",
+    "model = ONNXRunTime.load_inference(onnx_path)\n",
+    "\n",
+    "# Display the output classes we'll predict\n",
+    "println(\"The model predicts these flavor classes:\")\n",
+    "for class_name in config[\"output_names\"]\n",
+    "    println(\" - \", class_name)\n",
+    "end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "51c71156",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loaded 100000 events\n",
+      "Processing event #15\n",
+      "Loaded 65 reconstructed particles\n",
+      "Loaded 29 tracks\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Path to ROOT file with EDM4hep data\n",
+    "edm4hep_path = \"data/events_080263084.root\"\n",
+    "reader = RootIO.Reader(edm4hep_path)\n",
+    "\n",
+    "# Get event information\n",
+    "events = RootIO.get(reader, \"events\")\n",
+    "println(\"Loaded $(length(events)) events\")\n",
+    "\n",
+    "# Choose a specific event to analyze (event #15)\n",
+    "event_id = 15\n",
+    "evt = events[event_id]\n",
+    "println(\"Processing event #$event_id\")\n",
+    "\n",
+    "# Get reconstructed particles and tracks\n",
+    "recps = RootIO.get(reader, evt, \"ReconstructedParticles\")\n",
+    "tracks = RootIO.get(reader, evt, \"EFlowTrack_1\")\n",
+    "\n",
+    "# Get needed collections for feature extraction\n",
+    "bz = RootIO.get(reader, evt, \"magFieldBz\", register = false)[1]\n",
+    "trackdata = RootIO.get(reader, evt, \"EFlowTrack\")\n",
+    "trackerhits = RootIO.get(reader, evt, \"TrackerHits\")\n",
+    "gammadata = RootIO.get(reader, evt, \"EFlowPhoton\")\n",
+    "nhdata = RootIO.get(reader, evt, \"EFlowNeutralHadron\")\n",
+    "calohits = RootIO.get(reader, evt, \"CalorimeterHits\")\n",
+    "dNdx = RootIO.get(reader, evt, \"EFlowTrack_2\")\n",
+    "track_L = RootIO.get(reader, evt, \"EFlowTrack_L\", register = false)\n",
+    "\n",
+    "println(\"Loaded $(length(recps)) reconstructed particles\")\n",
+    "println(\"Loaded $(length(tracks)) tracks\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f96fe81b",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "2-element Vector{StructVector{ReconstructedParticle}}:\n",
+       " [ReconstructedParticle(#4, 0, 20.46865f0, (-19.487997, -1.9470071, 5.9491873), (0.0, 0.0, 0.0), -1.0f0, 0.00051099894f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[4], ReconstructedParticle#[], ParticleID#[4], #0, #0), ReconstructedParticle(#1, 0, 41.66906f0, (-38.420082, -8.888379, -13.460811), (0.0, 0.0, 0.0), 1.0f0, 0.105658375f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[1], ReconstructedParticle#[], ParticleID#[1], #0, #0), ReconstructedParticle(#30, 22, 5.97301f0, (-5.5111613, -1.3775525, -1.8456159), (0.0, 0.0, 0.0), 0.0f0, 0.00072206435f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[1], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0)]\n",
+       " [ReconstructedParticle(#49, 22, 0.038685087f0, (0.0016564546, 0.030539023, 0.023688823), (0.0, 0.0, 0.0), 0.0f0, -7.2657667f-6, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[20], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#50, 22, 0.07068448f0, (0.009334846, 0.04360084, 0.054846358), (0.0, 0.0, 0.0), 0.0f0, -1.211525f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[21], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#20, 0, 2.3770719f0, (0.18115185, 0.92913306, 2.1759794), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[20], ReconstructedParticle#[], ParticleID#[20], #0, #0), ReconstructedParticle(#52, 22, 0.49115527f0, (0.044869967, 0.21402507, 0.43978798), (0.0, 0.0, 0.0), 0.0f0, -6.145881f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[23], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#58, 22, 2.2526124f0, (-0.008166414, 0.70710534, 2.1387374), (0.0, 0.0, 0.0), 0.0f0, -0.00067013147f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[29], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#59, 22, 1.1003908f0, (-0.00017319276, 0.32061982, 1.0526457), (0.0, 0.0, 0.0), 0.0f0, -0.00021608449f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[30], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#54, 22, 1.757357f0, (0.08186126, 0.5793458, 1.6570942), (0.0, 0.0, 0.0), 0.0f0, -0.0007735479f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[25], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#56, 22, 5.843931f0, (0.19479804, 1.8359567, 5.5446234), (0.0, 0.0, 0.0), 0.0f0, -0.0012316033f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[27], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#25, 0, 2.5155842f0, (0.28472182, 0.63002336, 2.414682), (0.0, 0.0, 0.0), 1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[25], ReconstructedParticle#[], ParticleID#[25], #0, #0), ReconstructedParticle(#53, 22, 0.10352973f0, (0.032321043, 0.012602998, 0.09754445), (0.0, 0.0, 0.0), 0.0f0, 2.5103373f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[24], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0)  …  ReconstructedParticle(#36, 22, 1.2084657f0, (1.1462502, 0.09732889, 0.37017113), (0.0, 0.0, 0.0), 0.0f0, 9.1004724f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[7], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#7, 0, 3.8613276f0, (3.6221685, -0.1164991, 1.3254038), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[7], ReconstructedParticle#[], ParticleID#[7], #0, #0), ReconstructedParticle(#11, 0, 4.4815288f0, (4.0670757, -0.13165577, 1.8724797), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[11], ReconstructedParticle#[], ParticleID#[11], #0, #0), ReconstructedParticle(#28, 0, 0.5223787f0, (0.4566623, 0.0062101614, 0.21170896), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[28], ReconstructedParticle#[], ParticleID#[28], #0, #0), ReconstructedParticle(#41, 22, 0.3654142f0, (0.33327267, 0.007723306, 0.14965703), (0.0, 0.0, 0.0), 0.0f0, 3.5387056f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[12], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#10, 0, 1.0771525f0, (0.83581126, -0.55972606, 0.35903203), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[10], ReconstructedParticle#[], ParticleID#[10], #0, #0), ReconstructedParticle(#37, 22, 0.055122558f0, (0.046858095, -0.02121588, 0.0198167), (0.0, 0.0, 0.0), 0.0f0, 1.1097479f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[8], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#17, 0, 1.157116f0, (0.7993353, -0.37195656, 0.73630774), (0.0, 0.0, 0.0), -1.0f0, 0.13957039f0, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[], Track#[17], ReconstructedParticle#[], ParticleID#[17], #0, #0), ReconstructedParticle(#46, 22, 0.7516057f0, (0.59814185, -0.19748434, 0.41004556), (0.0, 0.0, 0.0), 0.0f0, -6.0317438f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[17], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0), ReconstructedParticle(#48, 22, 0.29473728f0, (0.23144196, -0.054464586, 0.1741789), (0.0, 0.0, 0.0), 0.0f0, 7.500272f-5, 0.0f0, Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], Cluster#[19], Track#[], ReconstructedParticle#[], ParticleID#[], #0, #0)]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# Cluster jets using the EEkt algorithm with R=2.0 and p=1.0\n",
+    "cs = jet_reconstruct(recps; p = 1.0, R = 2.0, algorithm = JetAlgorithm.EEKt)\n",
+    "\n",
+    "# Get 2 exclusive jets\n",
+    "jets = exclusive_jets(cs; njets=2, T=EEJet)\n",
+    "\n",
+    "# For each jet, get its constituent particles\n",
+    "constituent_indices = [constituent_indexes(jet, cs) for jet in jets]\n",
+    "jet_constituents = build_constituents_cluster(recps, constituent_indices)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "348a6c05",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "const JetConstituents = StructVector{ReconstructedParticle}\n",
+    "const JetConstituentsData = Vector{Float32}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "4f46c6c3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extracting features for flavor tagging...\n",
+      "Step 1: Feature extraction completed.\n",
+      "Step 2: Weaver setup completed.\n",
+      "Step 3: Input tensor preparation completed.\n",
+      "Running flavor tagging inference...\n",
+      "Step 4: Weights retrieval completed.\n",
+      "Jet scores:\n",
+      " - recojet_isG: 0.20045704\n",
+      " - recojet_isB: 0.7949233\n",
+      " - recojet_isQ: 0.0011744569\n",
+      " - recojet_isC: 0.0027449252\n",
+      " - recojet_isS: 0.00070025533\n"
+     ]
+    }
+   ],
+   "source": [
+    "println(\"Extracting features for flavor tagging...\")\n",
+    "feature_data = extract_features(\n",
+    "    jets, \n",
+    "    jet_constituents, \n",
+    "    tracks, \n",
+    "    bz, \n",
+    "    track_L, \n",
+    "    trackdata, \n",
+    "    trackerhits, \n",
+    "    gammadata, \n",
+    "    nhdata, \n",
+    "    calohits, \n",
+    "    dNdx\n",
+    ")\n",
+    "println(\"Step 1: Feature extraction completed.\")\n",
+    "\n",
+    "model, config = setup_weaver(\n",
+    "    onnx_path,\n",
+    "    json_path\n",
+    ")\n",
+    "\n",
+    "println(\"Step 2: Weaver setup completed.\")\n",
+    "\n",
+    "input_tensors = prepare_input_tensor(\n",
+    "    jet_constituents,\n",
+    "    jets,\n",
+    "    config,\n",
+    "    feature_data\n",
+    ")\n",
+    "\n",
+    "println(\"Step 3: Input tensor preparation completed.\")\n",
+    "\n",
+    "println(\"Running flavor tagging inference...\")\n",
+    "weights = get_weights(\n",
+    "    0,  # Thread slot\n",
+    "    feature_data,\n",
+    "    jets,\n",
+    "    jet_constituents,\n",
+    "    config,\n",
+    "    model\n",
+    ")\n",
+    "println(\"Step 4: Weights retrieval completed.\")\n",
+    "\n",
+    "jet_scores = Dict{String, Vector{Float32}}()\n",
+    "for (i, score_name) in enumerate(config[\"output_names\"])\n",
+    "    jet_scores[score_name] = get_weight(weights, i-1)\n",
+    "end\n",
+    "\n",
+    "println(\"Jet scores:\")\n",
+    "for (name, scores) in jet_scores\n",
+    "    println(\" - $name: $(scores[1])\")\n",
+    "end"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Julia 1.11.5",
+   "language": "julia",
+   "name": "julia-1.11"
+  },
+  "language_info": {
+   "file_extension": ".jl",
+   "mimetype": "application/julia",
+   "name": "julia",
+   "version": "1.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}