Update generative-proof-of-concept-CPU-preprocessing-in-memory.py

Add sqlite storage for Optuna.

Author: david-thrower

generative-proof-of-concept-CPU-preprocessing-in-memory.py

@@ -16,14 +16,16 @@
 EXPERIMENT_ITERATION = "0001"
 EXPERIMENT_NAME = "more-optimizations-br-254-single-machine"
 DATA_SET_NAME = "WEB-Bible-Genesis-40-context-681-SPL"
-
+EXPERIMENT_NAME = f"{EXPERIMENT_NAME}-{DATA_SET_NAME}-{EXPERIMENT_ITERATION}-a"
 
 N_TRIALS = 50
 
 
 mlflow.set_tracking_uri(uri=f"http://127.0.0.1:{MLFLOW_PORT}")
 
-mlflow.set_experiment(f"{EXPERIMENT_NAME}-{DATA_SET_NAME}-{EXPERIMENT_ITERATION}-a")
+
+
+mlflow.set_experiment(EXPERIMENT_NAME)
 
 
 
@@ -1443,7 +1445,8 @@ def main():
     n_trials = N_TRIALS
     # mlflow_parent = mlflow.start_run(run_name=os.getenv("MLFLOW_PARENT_RUN_NAME", "cerebros_poc_parent"), tags={"phase": "poc", "mode": "fast" if fast else "full"})
     sampler = optuna.samplers.TPESampler(multivariate=True, n_startup_trials=5)
-    study = optuna.create_study(direction="minimize", sampler=sampler)
+    storage_name = f"sqlite:///{EXPERIMENT_NAME}.db
+    study = optuna.create_study(direction="minimize", sampler=sampler, storage=storage_name)
     study.optimize(objective, n_trials=n_trials)
     # mlflow.log_param("n_trials", n_trials)
     # Log fixed (non-tunable) generation control param once at parent level