From 50cd905bfd9381bcfa6b77ae2c66af2ca69221b8 Mon Sep 17 00:00:00 2001
From: ramiz khaled <ramiz.khaled@inrae.fr>
Date: Tue, 14 Apr 2026 09:37:35 +0200
Subject: [PATCH] v0.1.6

---
 amn/_version.py |  2 +-
 amn/predict.py  |  5 +++++
 amn/train.py    | 15 +++++++++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/amn/_version.py b/amn/_version.py
index 1276d02..0a8da88 100644
--- a/amn/_version.py
+++ b/amn/_version.py
@@ -1 +1 @@
-__version__ = "0.1.5"
+__version__ = "0.1.6"
diff --git a/amn/predict.py b/amn/predict.py
index ae1681d..c36feda 100644
--- a/amn/predict.py
+++ b/amn/predict.py
@@ -60,6 +60,11 @@ def predict_amn(
         "biolog": {
             "cobra_file": get_default_model('biolog','build'),
             "metric": "r2_score"
+        },
+
+        "custom": {
+            "cobra_file": cobra_file,
+            "metric": metric
         }
     }
 
diff --git a/amn/train.py b/amn/train.py
index 39e07ff..f7395a8 100644
--- a/amn/train.py
+++ b/amn/train.py
@@ -59,6 +59,10 @@ def build_training_set(
 
         "biolog": {
             "cobra_file": get_default_model('biolog','build')
+        },
+
+        "custom": {
+            "cobra_file": cobra_file
         }
     }
     #if organism not in PRESETS:
@@ -147,6 +151,17 @@ def train_gr_prediction(
             "xfold": 5,
             "niter": 0,
             "batch_size": 100
+        },
+        
+        "custom": {
+            "trainingfile": trainingfile,
+            "cobraname_override": cobraname_override,
+            "n_hidden": n_hidden,
+            "hidden_dim": hidden_dim,
+            "epochs": epochs,
+            "xfold": xfold,
+            "niter": niter,
+            "batch_size": batch_size
         }
     }
     #if organism not in PRESETS: