Network graph not updating with js_on_change for a Select widget on Bokeh

Question

Following this example, I am trying to build a network graph using Bokeh where I can use a select widget to filter my data. My data will look something like this:

source  target  var1    var2
        a   c   1.0     0.0
        b   g   2.0     0.0
        c   e   3.0     0.0
        e   a   0.0     1.0
        f   h   0.0     2.0

And can be recreated using this code (note that I am pivoting the table so I will be filtering by columns and not by rows as per the example, but I am not sure if there is an easier way filtering by rows):

import pandas as pd
d = {'weight': [1, 2,3,1,2], 'var': ["var1","var1","var1","var2", "var2"], 'source': ["a", "b","c","e","f"], 'target': ["c","g","e","a","h"]}
df1 = pd.DataFrame(data=d)
df2 = df1.pivot( index= ["source","target"], values = "weight", columns = "var").reset_index()
df2 = df2.fillna(0)

Basically, I want to create the network graph where I can filter the columns (var1, var2) and they will become the weight attribute in my graph. (Filtering for when this weight value is greater than 0.)

To accomplish this I tried the following. But even though the graph renders, when I change the selected value nothing happens. I don't see any errors in the console either. I am not sure what I am doing wrong, probably something in the JS call because I am new to this, but I am trying to reproduce the example as closely as possible and still not sure where im going wrong. Please help!

from bokeh.plotting import from_networkx
from bokeh.plotting import figure, output_file, show
from bokeh.models import Plot, Range1d, MultiLine, Circle, HoverTool,NodesAndLinkedEdges,EdgesAndLinkedNodes, TapTool, BoxSelectTool,ColumnDataSource
from bokeh.models import CustomJS, ColumnDataSource, Select, Column
import networkx as nx
import networkx

HOVER_TOOLTIPS = [("Search Term", "@index")]

title = "my title"
plot = figure(tooltips = HOVER_TOOLTIPS,
              tools="pan,wheel_zoom", 
              active_scroll='wheel_zoom',
              x_range=Range1d(-10.1, 10.1), 
              y_range=Range1d(-10.1, 10.1), 
              title=title, plot_width=1000
             )

category_default = "var1"
unique_categories = ["var1","var2"]

subset = df2
subset = subset.fillna(0)
subset_data = subset[["var1","var2"]].to_dict("list") 
subset = subset.set_index(["source","target"])

source = ColumnDataSource({
    
        "source": subset.index,
        "weight": subset[category_default]   


    })

select = Select(title='Category Selection', value=category_default, options=unique_categories)

callback = CustomJS(
        args={"subset_data": subset_data, "source": source},
    
        code= """ 
            source.data["weight"] = subset_data[cb_obj.value];
            source.change.emit();          
        """)

select.js_on_change("value", callback)

b = pd.DataFrame(source.data["source"].to_frame(index=False))
b["weight"] = source.data["weight"].values

G = nx.from_pandas_edgelist(b[b["weight"] >0], edge_attr = "weight")

network_graph = from_networkx(G, networkx.spring_layout, scale=10, center=(0, 0), )
network_graph.edge_renderer.glyph = MultiLine(line_alpha=0.5, line_width="weight" )

plot.renderers.append(network_graph)
show(Column(plot, select))