memory management

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import pandas as pd
 import traceback
+import gc
+import torch
 
 from cognitive_mapping_probe.orchestrator_seismograph import run_seismic_analysis
 from cognitive_mapping_probe.auto_experiment import run_auto_suite, get_curated_experiments
@@ -9,6 +11,14 @@ from cognitive_mapping_probe.utils import dbg
 
 theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue").set(body_background_fill="#f0f4f9", block_background_fill="white")
 
+def cleanup_memory():
+    """Eine zentrale Funktion zum Aufräumen des Speichers nach einem Lauf."""
+    dbg("Cleaning up memory...")
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    dbg("Memory cleanup complete.")
+
 def run_single_analysis_display(*args, progress=gr.Progress(track_tqdm=True)):
     """Wrapper für ein einzelnes manuelles Experiment."""
     try:
@@ -17,16 +27,26 @@ def run_single_analysis_display(*args, progress=gr.Progress(track_tqdm=True)):
         deltas = results.get("state_deltas", [])
         df = pd.DataFrame({"Internal Step": range(len(deltas)), "State Change (Delta)": deltas})
         stats_md = f"### Statistical Signature\n- **Mean Delta:** {stats.get('mean_delta', 0):.4f}\n- **Std Dev Delta:** {stats.get('std_delta', 0):.4f}\n- **Max Delta:** {stats.get('max_delta', 0):.4f}\n"
+
+        # WICHTIG: Speicher aufräumen, BEVOR die Ergebnisse an Gradio zurückgegeben werden.
+        cleanup_memory()
+
         return f"{results.get('verdict', 'Error')}\n\n{stats_md}", df, results
     except Exception:
+        cleanup_memory()
         return f"### ❌ Analysis Failed\n```\n{traceback.format_exc()}\n```", pd.DataFrame(), {}
 
 def run_auto_suite_display(model_id, num_steps, seed, experiment_name, progress=gr.Progress(track_tqdm=True)):
     """Wrapper für die automatisierte Experiment-Suite mit Visualisierung."""
     try:
         summary_df, plot_df, all_results = run_auto_suite(model_id, int(num_steps), int(seed), experiment_name, progress)
+
+        # WICHTIG: Speicher auch hier aufräumen.
+        cleanup_memory()
+
         return summary_df, plot_df, all_results
     except Exception:
+        cleanup_memory()
         return pd.DataFrame(), pd.DataFrame(), f"### ❌ Auto-Experiment Failed\n```\n{traceback.format_exc()}\n```"
 
 with gr.Blocks(theme=theme, title="Cognitive Seismograph 2.2") as demo:
@@ -49,7 +69,6 @@ with gr.Blocks(theme=theme, title="Cognitive Seismograph 2.2") as demo:
                 with gr.Column(scale=2):
                     gr.Markdown("### Single Run Results")
                     manual_verdict = gr.Markdown("Die Analyse erscheint hier.")
-                    # KORREKTUR: `interactive=True` für Legende hinzugefügt
                     manual_plot = gr.LinePlot(x="Internal Step", y="State Change (Delta)", title="Internal State Dynamics", show_label=True, height=400, interactive=True)
                     with gr.Accordion("Raw JSON Output", open=False):
                         manual_raw_json = gr.JSON()
@@ -72,7 +91,6 @@ with gr.Blocks(theme=theme, title="Cognitive Seismograph 2.2") as demo:
                     auto_run_btn = gr.Button("Run Curated Auto-Experiment", variant="primary")
                 with gr.Column(scale=2):
                     gr.Markdown("### Suite Results Summary")
-                    # KORREKTUR: `interactive=True` für Legende hinzugefügt
                     auto_plot_output = gr.LinePlot(
                         x="Step", y="Delta", color="Experiment",
                         title="Comparative Cognitive Dynamics",

cognitive_mapping_probe/auto_experiment.py

@@ -1,4 +1,6 @@
 import pandas as pd
+import torch
+import gc
 from typing import Dict, List, Tuple
 
 from .llm_iface import get_or_load_model
@@ -8,7 +10,6 @@ from .utils import dbg
 def get_curated_experiments() -> Dict[str, List[Dict]]:
     """
     Definiert die vordefinierten, wissenschaftlichen Experiment-Protokolle.
-    ERWEITERT um zusätzliche, aussagekräftige Experimente.
     """
     experiments = {
         "Calm vs. Chaos": [
@@ -61,6 +62,7 @@ def run_auto_suite(
     summary_data = []
     plot_data_frames = []
 
+    # Lade das Modell einmal zu Beginn der Suite
     llm = get_or_load_model(model_id, seed)
 
     total_runs = len(protocol)
@@ -69,40 +71,29 @@ def run_auto_suite(
         dbg(f"--- Running Auto-Experiment: '{label}' ({i+1}/{total_runs}) ---")
 
         results = run_seismic_analysis(
-            model_id=model_id,
-            prompt_type=run_spec["prompt_type"],
-            seed=seed,
-            num_steps=num_steps,
-            concept_to_inject=run_spec["concept"],
-            injection_strength=run_spec["strength"],
+            model_id=model_id, prompt_type=run_spec["prompt_type"],
+            seed=seed, num_steps=num_steps,
+            concept_to_inject=run_spec["concept"], injection_strength=run_spec["strength"],
             progress_callback=progress_callback,
-            llm_instance=llm
+            llm_instance=llm # Wiederverwende die geladene LLM-Instanz
         )
 
         all_results[label] = results
         stats = results.get("stats", {})
 
-        summary_data.append({
-            "Experiment": label,
-            "Mean Delta": stats.get("mean_delta"),
-            "Std Dev Delta": stats.get("std_delta"),
-            "Max Delta": stats.get("max_delta"),
-        })
+        summary_data.append({ "Experiment": label, "Mean Delta": stats.get("mean_delta"), "Std Dev Delta": stats.get("std_delta"), "Max Delta": stats.get("max_delta"), })
 
         deltas = results.get("state_deltas", [])
-
-        # Verwende das Label direkt für die Legende. Der vorherige Fix war nicht nötig,
-        # solange die Labels einzigartig pro Experiment sind.
-        df = pd.DataFrame({
-            "Step": range(len(deltas)),
-            "Delta": deltas,
-            "Experiment": label
-        })
+        df = pd.DataFrame({ "Step": range(len(deltas)), "Delta": deltas, "Experiment": f"{i}: {label}" })
         plot_data_frames.append(df)
 
     summary_df = pd.DataFrame(summary_data)
     plot_df = pd.concat(plot_data_frames, ignore_index=True) if plot_data_frames else pd.DataFrame()
 
+    # WICHTIG: Explizites Aufräumen am Ende der gesamten Suite
     del llm
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
     return summary_df, plot_df, all_results

cognitive_mapping_probe/orchestrator_seismograph.py

@@ -1,5 +1,6 @@
 import torch
 import numpy as np
+import gc
 from typing import Dict, Any, Optional
 
 from .llm_iface import get_or_load_model
@@ -15,39 +16,35 @@ def run_seismic_analysis(
     concept_to_inject: str,
     injection_strength: float,
     progress_callback,
-    llm_instance: Optional[Any] = None # Ermöglicht Wiederverwendung des Modells
+    llm_instance: Optional[Any] = None
 ) -> Dict[str, Any]:
     """
-    Orchestriert eine einzelne seismische Analyse. Kann optional eine bestehende
-    LLM-Instanz wiederverwenden, um das Neuladen in automatisierten Suiten zu beschleunigen.
+    Orchestriert eine einzelne seismische Analyse. Stellt sicher, dass das Modell
+    nur dann entladen wird, wenn es auch hier geladen wurde.
     """
-    # Lade das Modell nur, wenn keine Instanz übergeben wurde
+    local_llm_instance = False
     if llm_instance is None:
         progress_callback(0.1, desc="Loading model...")
         llm = get_or_load_model(model_id, seed)
-        created_llm = True
+        local_llm_instance = True
     else:
         llm = llm_instance
-        llm.set_all_seeds(seed) # Setze den Seed für diesen spezifischen Lauf
-        created_llm = False
+        llm.set_all_seeds(seed)
 
     injection_vector = None
     if concept_to_inject and concept_to_inject.strip():
-        if not created_llm: progress_callback(0.2, desc=f"Vectorizing '{concept_to_inject}'...")
+        if not local_llm_instance: progress_callback(0.2, desc=f"Vectorizing '{concept_to_inject}'...")
         injection_vector = get_concept_vector(llm, concept_to_inject.strip())
 
-    if not created_llm: progress_callback(0.3, desc=f"Recording dynamics...")
+    if not local_llm_instance: progress_callback(0.3, desc=f"Recording dynamics...")
 
     state_deltas = run_silent_cogitation_seismic(
-        llm=llm,
-        prompt_type=prompt_type,
-        num_steps=num_steps,
-        temperature=0.1,
-        injection_vector=injection_vector,
-        injection_strength=injection_strength
+        llm=llm, prompt_type=prompt_type,
+        num_steps=num_steps, temperature=0.1,
+        injection_vector=injection_vector, injection_strength=injection_strength
     )
 
-    if not created_llm: progress_callback(0.9, desc="Analyzing...")
+    if not local_llm_instance: progress_callback(0.9, desc="Analyzing...")
 
     if state_deltas:
         deltas_np = np.array(state_deltas)
@@ -60,9 +57,14 @@ def run_seismic_analysis(
 
     results = { "verdict": verdict, "stats": stats, "state_deltas": state_deltas }
 
-    # Gib das Modell nur frei, wenn es in dieser Funktion erstellt wurde
-    if created_llm:
+    # WICHTIG: Gib das Modell und den Speicher nur frei, wenn es in dieser
+    # Funktion auch erstellt wurde. Ansonsten ist die übergeordnete Funktion
+    # (z.B. `run_auto_suite`) für das Speichermanagement verantwortlich.
+    if local_llm_instance:
         del llm
-        if torch.cuda.is_available(): torch.cuda.empty_cache()
+        del injection_vector
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
     return results