fix: added function production_instantiate

Author: hack-hard

src/models.jl

@@ -4,14 +4,14 @@ function select_and_preprocess((point, atoms); cutoff_radius)
     select_and_preprocess(point, atoms; cutoff_radius)
 end
 function select_and_preprocess(
-        point::Batch, atoms::AnnotedKDTree{Sphere{T}}; cutoff_radius) where {T}
+    point::Batch, atoms::AnnotedKDTree{Sphere{T}}; cutoff_radius) where {T}
     neighboord = Folds.map(point.field) do point
         select_neighboord(point,
             atoms;
             cutoff_radius)::StructVector{
-            Sphere{T}, @NamedTuple{center::Vector{Point3{T}}, r::Vector{T}}, Int64}
+            Sphere{T},@NamedTuple{center::Vector{Point3{T}}, r::Vector{T}},Int64}
     end |> Batch{Vector{StructVector{
-        Sphere{T}, @NamedTuple{center::Vector{Point3{T}}, r::Vector{T}}, Int64}}}
+        Sphere{T},@NamedTuple{center::Vector{Point3{T}}, r::Vector{T}},Int64}}}
     preprocessing((point, neighboord))
 end
 
@@ -24,9 +24,9 @@ function evaluate_if_atoms_in_neighboord(layer, arg::AbstractArray, ps, st; zero
 end
 
 function general_angular_dense(main_chain, secondary_chain; name::String,
-        van_der_waals_channel = false, on_gpu = true, cutoff_radius::Float32 = 3.0f0)
+    van_der_waals_channel=false, on_gpu=true, cutoff_radius::Float32=3.0f0)
     main_chain = DeepSet(Chain(
-        symetrise(; cutoff_radius, device = on_gpu ? gpu_device() : identity),
+        symetrise(; cutoff_radius, device=on_gpu ? gpu_device() : identity),
         main_chain
     ))
     function add_van_der_waals_channel(main_chain)
@@ -46,7 +46,7 @@ end
 	`tiny_angular_dense` is a function that generate a lux model.
 
 """
-function tiny_angular_dense(; van_der_waals_channel = false, kargs...)
+function tiny_angular_dense(; van_der_waals_channel=false, kargs...)
     general_angular_dense(
         Parallel(.*,
             Chain(Dense(6 => 7, elu),
@@ -57,12 +57,12 @@ function tiny_angular_dense(; van_der_waals_channel = false, kargs...)
             BatchNorm(4 + van_der_waals_channel),
             Dense(4 + van_der_waals_channel => 6, elu),
             Dense(6 => 1, sigmoid_fast));
-        name = "tiny_angular_dense_" *
-               (van_der_waals_channel ? "v" : ""),
+        name="tiny_angular_dense_" *
+             (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel, kargs...)
 end
 
-function light_angular_dense(; van_der_waals_channel = false, kargs...)
+function light_angular_dense(; van_der_waals_channel=false, kargs...)
     general_angular_dense(
         Parallel(.*,
             Chain(Dense(6 => 10, elu),
@@ -73,13 +73,13 @@ function light_angular_dense(; van_der_waals_channel = false, kargs...)
             BatchNorm(5 + van_der_waals_channel),
             Dense(5 + van_der_waals_channel => 10, elu),
             Dense(10 => 1, sigmoid_fast));
-        name = "light_angular_dense_" *
-               (van_der_waals_channel ? "v" : ""),
+        name="light_angular_dense_" *
+             (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel, kargs...)
 end
 
 function medium_angular_dense(;
-        van_der_waals_channel = false, kargs...)
+    van_der_waals_channel=false, kargs...)
     general_angular_dense(
         Parallel(.*,
             Chain(Dense(6 => 15, elu),
@@ -88,17 +88,17 @@ function medium_angular_dense(;
         ),
         Chain(
             BatchNorm(10 + van_der_waals_channel),
-            Dense(10 + van_der_waals_channel => 5; use_bias = false),
+            Dense(10 + van_der_waals_channel => 5; use_bias=false),
             Dense(5 => 10, elu),
             Dense(10 => 1, sigmoid_fast));
-        name = "medium_angular_dense_" *
-               (van_der_waals_channel ? "v" : ""),
+        name="medium_angular_dense_" *
+             (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel,
         kargs...)
 end
 function drop_preprocessing(x::Chain)
     if typeof(x[1]) <: PreprocessingLayer
-        Chain(NoOpLayer(), map(i -> x[i], 2:length(x))..., disable_optimizations = true)
+        Chain(NoOpLayer(), map(i -> x[i], 2:length(x))..., disable_optimizations=true)
     else
         x
     end
@@ -115,6 +115,11 @@ struct SerializedModel
     model::Partial
     weights::NamedTuple
 end
+production_instantiate((; model, weights)::SerializedModel) = Lux.StatefulLuxLayer(
+    model(), 
+	weights, 
+	Lux.initialparameters(MersenneTwister(42), model.model()) |> Lux.testmode)
+
 function get_cutoff_radius(x::Lux.AbstractExplicitLayer)
     get_preprocessing(x).fun.kargs[:cutoff_radius]
 end