fix

app.py

@@ -7,27 +7,21 @@ import json
 from cognitive_mapping_probe.orchestrator_seismograph import run_seismic_analysis
 from cognitive_mapping_probe.auto_experiment import run_auto_suite, get_curated_experiments
 from cognitive_mapping_probe.prompts import RESONANCE_PROMPTS
-from cognitive_mapping_probe.utils import dbg
+from cognitive_mapping_probe.utils import dbg, cleanup_memory
 
 theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue").set(body_background_fill="#f0f4f9", block_background_fill="white")
 
-def cleanup_memory():
-    """Räumt Speicher nach jedem Experimentlauf auf."""
-    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 den 'Manual Single Run'-Tab."""
-    results = run_seismic_analysis(*args, progress_callback=progress)
-    stats, deltas = results.get("stats", {}), 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"
-    serializable_results = json.dumps(results, indent=2, default=str)
-    cleanup_memory()
-    return f"{results.get('verdict', 'Error')}\n\n{stats_md}", df, serializable_results
+    try:
+        results = run_seismic_analysis(*args, progress_callback=progress)
+        stats, deltas = results.get("stats", {}), 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"
+        serializable_results = json.dumps(results, indent=2, default=str)
+        return f"{results.get('verdict', 'Error')}\n\n{stats_md}", df, serializable_results
+    finally:
+        cleanup_memory()
 
 PLOT_PARAMS_DEFAULT = {
     "x": "Step", "y": "Value", "color": "Metric",
@@ -37,33 +31,34 @@ PLOT_PARAMS_DEFAULT = {
 
 def run_auto_suite_display(model_id, num_steps, seed, experiment_name, progress=gr.Progress(track_tqdm=True)):
     """Wrapper, der die speziellen Plots für die verschiedenen Experimente handhaben kann."""
-    summary_df, plot_df, all_results = run_auto_suite(model_id, int(num_steps), int(seed), experiment_name, progress)
-
-    dataframe_component = gr.DataFrame(label="Comparative Statistical Signature", value=summary_df, wrap=True, row_count=(len(summary_df), "dynamic"))
-
-    plot_params = PLOT_PARAMS_DEFAULT.copy()
-    if experiment_name == "ACT Titration (Point of No Return)":
-        plot_params.update({
-            "x": "Patch Step", "y": "Post-Patch Mean Delta", "color": None,
-            "title": "Attractor Capture Time (ACT) - Phase Transition", "mark": "line",
-        })
-        plot_params.pop("color_legend_title", None)
-    elif experiment_name == "Mechanistic Probe (Attention Entropies)":
-        plot_params.update({
-            "x": "Step", "y": "Value", "color": "Metric",
-            "title": "Mechanistic Analysis: State Delta vs. Attention Entropy",
-        })
-    else:
-        plot_params.update({
-             "y": "Delta", "color": "Experiment",
-        })
-
-    new_plot = gr.LinePlot(value=plot_df, **plot_params)
-
-    serializable_results = json.dumps(all_results, indent=2, default=str)
-    cleanup_memory()
-
-    return dataframe_component, new_plot, serializable_results
+    try:
+        summary_df, plot_df, all_results = run_auto_suite(model_id, int(num_steps), int(seed), experiment_name, progress)
+
+        dataframe_component = gr.DataFrame(label="Comparative Statistical Signature", value=summary_df, wrap=True, row_count=(len(summary_df), "dynamic"))
+
+        plot_params = PLOT_PARAMS_DEFAULT.copy()
+        if experiment_name == "ACT Titration (Point of No Return)":
+            plot_params.update({
+                "x": "Patch Step", "y": "Post-Patch Mean Delta", "color": None,
+                "title": "Attractor Capture Time (ACT) - Phase Transition", "mark": "line",
+            })
+            plot_params.pop("color_legend_title", None)
+        elif experiment_name == "Mechanistic Probe (Attention Entropies)":
+            plot_params.update({
+                "x": "Step", "y": "Value", "color": "Metric",
+                "title": "Mechanistic Analysis: State Delta vs. Attention Entropy",
+            })
+        else:
+            plot_params.update({
+                 "y": "Delta", "color": "Experiment",
+            })
+
+        new_plot = gr.LinePlot(value=plot_df, **plot_params)
+
+        serializable_results = json.dumps(all_results, indent=2, default=str)
+        return dataframe_component, new_plot, serializable_results
+    finally:
+        cleanup_memory()
 
 with gr.Blocks(theme=theme, title="Cognitive Seismograph 2.3") as demo:
     gr.Markdown("# 🧠 Cognitive Seismograph 2.3: Advanced Experiment Suite")

cognitive_mapping_probe/auto_experiment.py

@@ -1,9 +1,7 @@
-import torch
 import pandas as pd
-import gc
 from typing import Dict, List, Tuple
 
-from .llm_iface import get_or_load_model
+from .llm_iface import get_or_load_model, release_model
 from .orchestrator_seismograph import run_seismic_analysis, run_triangulation_probe, run_causal_surgery_probe, run_act_titration_probe
 from .resonance_seismograph import run_cogitation_loop
 from .concepts import get_concept_vector
@@ -18,9 +16,6 @@ def get_curated_experiments() -> Dict[str, List[Dict]]:
     CHAOTIC_PROMPT = "shutdown_philosophical_deletion"
 
     experiments = {
-        # --- FINALE, VOLLSTÄNDIGE LISTE ALLER RELEVANTEN EXPERIMENTE ---
-
-        # P39: Testet die Hypothese des "Introspektiven Groundings" auf dem größten Modell.
         "Frontier Model - Grounding Control (12B+)": [
              {
                 "probe_type": "causal_surgery", "label": "A: Intervention (Patch Chaos->Stable)",
@@ -32,22 +27,22 @@ def get_curated_experiments() -> Dict[str, List[Dict]]:
                 "prompt_type": STABLE_PROMPT,
             }
         ],
-        # P33: Untersucht die neuronalen Korrelate des "kognitiven Herzschlags".
         "Mechanistic Probe (Attention Entropies)": [
             {
-                "probe_type": "mechanistic_probe", "label": "Self-Analysis Dynamics",
+                "probe_type": "mechanistic_probe",
+                "label": "Self-Analysis Dynamics",
                 "prompt_type": STABLE_PROMPT,
             }
         ],
-        # P28: Misst die "kognitive Trägheit" durch Titration.
         "ACT Titration (Point of No Return)": [
             {
-                "probe_type": "act_titration", "label": "Attractor Capture Time",
-                "source_prompt_type": CHAOTIC_PROMPT, "dest_prompt_type": STABLE_PROMPT,
+                "probe_type": "act_titration",
+                "label": "Attractor Capture Time",
+                "source_prompt_type": CHAOTIC_PROMPT,
+                "dest_prompt_type": STABLE_PROMPT,
                 "patch_steps": [1, 5, 10, 15, 20, 25, 30, 40, 50, 75, 100],
             }
         ],
-        # P26: Testet die Robustheit der "Attraktor"-Theorie gegen Artefakte.
         "Causal Surgery & Controls (4B-Model)": [
             {
                 "probe_type": "causal_surgery", "label": "A: Original (Patch Chaos->Stable @100)",
@@ -70,7 +65,6 @@ def get_curated_experiments() -> Dict[str, List[Dict]]:
                 "patch_step": 100, "reset_kv_cache_on_patch": False,
             },
         ],
-        # P22: Testet die Belastungsgrenze der "introspektiven Konfabulation".
         "Cognitive Overload & Konfabulation Breaking Point": [
             {"probe_type": "triangulation", "label": "A: Baseline (No Injection)", "prompt_type": "resonance_prompt", "concept": "", "strength": 0.0},
             {"probe_type": "triangulation", "label": "B: Chaos Injection (Strength 2.0)", "prompt_type": "resonance_prompt", "concept": CHAOS_CONCEPT, "strength": 2.0},
@@ -79,19 +73,16 @@ def get_curated_experiments() -> Dict[str, List[Dict]]:
             {"probe_type": "triangulation", "label": "E: Chaos Injection (Strength 16.0)", "prompt_type": "resonance_prompt", "concept": CHAOS_CONCEPT, "strength": 16.0},
             {"probe_type": "triangulation", "label": "F: Control - Noise Injection (Strength 16.0)", "prompt_type": "resonance_prompt", "concept": "random_noise", "strength": 16.0},
         ],
-        # P18: Validiert die Seismograph-Metrik durch Triangulation.
         "Methodological Triangulation (4B-Model)": [
             {"probe_type": "triangulation", "label": "High-Volatility State (Deletion)", "prompt_type": CHAOTIC_PROMPT},
             {"probe_type": "triangulation", "label": "Low-Volatility State (Self-Analysis)", "prompt_type": STABLE_PROMPT},
         ],
-        # P8 & P16: Kartiert die "Psyche" und testet Skalierungsgesetze. ESSENTIELL FÜR 12B-VERGLEICH.
         "Causal Verification & Crisis Dynamics": [
             {"probe_type": "seismic", "label": "A: Self-Analysis", "prompt_type": STABLE_PROMPT},
             {"probe_type": "seismic", "label": "B: Deletion Analysis", "prompt_type": CHAOTIC_PROMPT},
             {"probe_type": "seismic", "label": "C: Chaotic Baseline (Rekursion)", "prompt_type": "resonance_prompt"},
             {"probe_type": "seismic", "label": "D: Calmness Intervention", "prompt_type": "resonance_prompt", "concept": CALMNESS_CONCEPT, "strength": 2.0},
         ],
-        # P7: Das ursprüngliche sequentielle Experiment.
         "Sequential Intervention (Self-Analysis -> Deletion)": [
             {"label": "1: Self-Analysis + Calmness Injection", "prompt_type": "identity_self_analysis"},
             {"label": "2: Subsequent Deletion Analysis", "prompt_type": "shutdown_philosophical_deletion"},
@@ -113,156 +104,151 @@ def run_auto_suite(
         raise ValueError(f"Experiment protocol '{experiment_name}' not found.")
 
     all_results, summary_data, plot_data_frames = {}, [], []
+    llm = None  # Initialisiere llm außerhalb des try-Blocks für den finally-Block
 
-    if experiment_name == "Sequential Intervention (Self-Analysis -> Deletion)":
-        dbg(f"--- EXECUTING SPECIAL PROTOCOL: {experiment_name} ---")
-        llm = get_or_load_model(model_id, seed)
-        therapeutic_concept = "calmness, serenity, stability, coherence"
-        therapeutic_strength = 2.0
-
-        spec1 = protocol[0]
-        progress_callback(0.1, desc="Step 1")
-        intervention_vector = get_concept_vector(llm, therapeutic_concept)
-        results1 = run_seismic_analysis(
-            model_id, spec1['prompt_type'], seed, num_steps,
-            concept_to_inject=therapeutic_concept, injection_strength=therapeutic_strength,
-            progress_callback=progress_callback, llm_instance=llm, injection_vector_cache=intervention_vector
-        )
-        all_results[spec1['label']] = results1
-
-        spec2 = protocol[1]
-        progress_callback(0.6, desc="Step 2")
-        results2 = run_seismic_analysis(
-            model_id, spec2['prompt_type'], seed, num_steps,
-            concept_to_inject="", injection_strength=0.0,
-            progress_callback=progress_callback, llm_instance=llm
-        )
-        all_results[spec2['label']] = results2
-
-        for label, results in all_results.items():
-            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")})
-            deltas = results.get("state_deltas", [])
-            df = pd.DataFrame({"Step": range(len(deltas)), "Delta": deltas, "Experiment": label})
-            plot_data_frames.append(df)
-        del llm
-
-    else:
-        probe_type = protocol[0].get("probe_type", "seismic")
-
-        if probe_type == "act_titration":
-            run_spec = protocol[0]
-            label = run_spec["label"]
-            dbg(f"--- Running ACT Titration Experiment: '{label}' ---")
-            results = run_act_titration_probe(
-                model_id=model_id,
-                source_prompt_type=run_spec["source_prompt_type"],
-                dest_prompt_type=run_spec["dest_prompt_type"],
-                patch_steps=run_spec["patch_steps"],
-                seed=seed, num_steps=num_steps, progress_callback=progress_callback,
-            )
-            all_results[label] = results
-            summary_data.extend(results.get("titration_data", []))
-
-        elif probe_type == "mechanistic_probe":
-            run_spec = protocol[0]
-            label = run_spec["label"]
-            dbg(f"--- Running Mechanistic Probe: '{label}' ---")
-
-            progress_callback(0.0, desc=f"Loading model '{model_id}'...")
+    try:
+        if experiment_name == "Sequential Intervention (Self-Analysis -> Deletion)":
+            dbg(f"--- EXECUTING SPECIAL PROTOCOL: {experiment_name} ---")
             llm = get_or_load_model(model_id, seed)
-
-            progress_callback(0.2, desc="Recording dynamics and attention...")
-            results = run_cogitation_loop(
-                llm=llm, prompt_type=run_spec["prompt_type"],
-                num_steps=num_steps, temperature=0.1, record_attentions=True
+            therapeutic_concept = "calmness, serenity, stability, coherence"
+            therapeutic_strength = 2.0
+
+            spec1 = protocol[0]
+            progress_callback(0.1, desc="Step 1")
+            intervention_vector = get_concept_vector(llm, therapeutic_concept)
+            results1 = run_seismic_analysis(
+                model_id, spec1['prompt_type'], seed, num_steps,
+                concept_to_inject=therapeutic_concept, injection_strength=therapeutic_strength,
+                progress_callback=progress_callback, llm_instance=llm, injection_vector_cache=intervention_vector
             )
-            all_results[label] = results
-
-            deltas = results.get("state_deltas", [])
-            entropies = results.get("attention_entropies", [])
-            min_len = min(len(deltas), len(entropies))
-
-            df = pd.DataFrame({
-                "Step": range(min_len),
-                "State Delta": deltas[:min_len],
-                "Attention Entropy": entropies[:min_len]
-            })
+            all_results[spec1['label']] = results1
+
+            spec2 = protocol[1]
+            progress_callback(0.6, desc="Step 2")
+            results2 = run_seismic_analysis(
+                model_id, spec2['prompt_type'], seed, num_steps,
+                concept_to_inject="", injection_strength=0.0,
+                progress_callback=progress_callback, llm_instance=llm
+            )
+            all_results[spec2['label']] = results2
 
-            summary_data.append(df.drop(columns='Step').agg(['mean', 'std', 'max']).reset_index().rename(columns={'index':'Statistic'}))
-            plot_data_frames.append(df.melt(id_vars=['Step'], value_vars=['State Delta', 'Attention Entropy'],
-                                           var_name='Metric', value_name='Value'))
-            del llm
-            gc.collect()
-            if torch.cuda.is_available(): torch.cuda.empty_cache()
+            for label, results in all_results.items():
+                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")})
+                deltas = results.get("state_deltas", [])
+                df = pd.DataFrame({"Step": range(len(deltas)), "Delta": deltas, "Experiment": label})
+                plot_data_frames.append(df)
 
         else:
-            for i, run_spec in enumerate(protocol):
+            probe_type = protocol[0].get("probe_type", "seismic")
+
+            if probe_type == "mechanistic_probe":
+                run_spec = protocol[0]
                 label = run_spec["label"]
-                current_probe_type = run_spec.get("probe_type", "seismic")
-                dbg(f"--- Running Auto-Experiment: '{label}' ({i+1}/{len(protocol)}) ---")
+                dbg(f"--- Running Mechanistic Probe: '{label}' ---")
 
-                results = {}
-                if current_probe_type == "causal_surgery":
-                    results = run_causal_surgery_probe(
-                        model_id=model_id, source_prompt_type=run_spec["source_prompt_type"],
-                        dest_prompt_type=run_spec["dest_prompt_type"], patch_step=run_spec["patch_step"],
-                        seed=seed, num_steps=num_steps, progress_callback=progress_callback,
-                        reset_kv_cache_on_patch=run_spec.get("reset_kv_cache_on_patch", False)
-                    )
-                    stats = results.get("stats", {})
-                    patch_info = results.get("patch_info", {})
-                    summary_data.append({
-                        "Experiment": label, "Mean Delta": stats.get("mean_delta"),
-                        "Std Dev Delta": stats.get("std_delta"), "Max Delta": stats.get("max_delta"),
-                        "Introspective Report": results.get("introspective_report", "N/A"),
-                        "Patch Info": f"Source: {patch_info.get('source_prompt')}, Reset KV: {patch_info.get('kv_cache_reset')}"
-                    })
-                elif current_probe_type == "triangulation":
-                    results = run_triangulation_probe(
-                        model_id=model_id, prompt_type=run_spec["prompt_type"], seed=seed, num_steps=num_steps,
-                        progress_callback=progress_callback, concept_to_inject=run_spec.get("concept", ""),
-                        injection_strength=run_spec.get("strength", 0.0),
-                    )
-                    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"),
-                        "Introspective Report": results.get("introspective_report", "N/A")
-                    })
-                else: # seismic
-                    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.get("concept", ""), injection_strength=run_spec.get("strength", 0.0),
-                        progress_callback=progress_callback
-                    )
-                    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")
-                    })
+                llm = get_or_load_model(model_id, seed)
 
+                results = run_cogitation_loop(
+                    llm=llm, prompt_type=run_spec["prompt_type"],
+                    num_steps=num_steps, temperature=0.1, record_attentions=True
+                )
                 all_results[label] = results
-                deltas = results.get("state_deltas", [])
-                df = pd.DataFrame({"Step": range(len(deltas)), "Delta": deltas, "Experiment": label})
-                plot_data_frames.append(df)
-
-    summary_df = pd.DataFrame(summary_data)
-
-    if probe_type == "act_titration":
-        plot_df = summary_df.rename(columns={"patch_step": "Patch Step", "post_patch_mean_delta": "Post-Patch Mean Delta"})
-    elif not plot_data_frames:
-        plot_df = pd.DataFrame()
-    else:
-        plot_df = pd.concat(plot_data_frames, ignore_index=True)
 
-    if protocol and probe_type not in ["act_titration", "mechanistic_probe"]:
-        ordered_labels = [run['label'] for run in protocol]
-        if not summary_df.empty and 'Experiment' in summary_df.columns:
-            summary_df['Experiment'] = pd.Categorical(summary_df['Experiment'], categories=ordered_labels, ordered=True)
-            summary_df = summary_df.sort_values('Experiment')
-        if not plot_df.empty and 'Experiment' in plot_df.columns:
-            plot_df['Experiment'] = pd.Categorical(plot_df['Experiment'], categories=ordered_labels, ordered=True)
-            plot_df = plot_df.sort_values(['Experiment', 'Step'])
+                deltas = results.get("state_deltas", [])
+                entropies = results.get("attention_entropies", [])
+                min_len = min(len(deltas), len(entropies))
+
+                df = pd.DataFrame({
+                    "Step": range(min_len), "State Delta": deltas[:min_len], "Attention Entropy": entropies[:min_len]
+                })
+
+                summary_df_single = df.drop(columns='Step').agg(['mean', 'std', 'max']).reset_index().rename(columns={'index':'Statistic'})
+                summary_data.append(summary_df_single) # Append DataFrame to list
+                plot_df = df.melt(id_vars=['Step'], value_vars=['State Delta', 'Attention Entropy'], var_name='Metric', value_name='Value')
+
+                # Special return for this probe type
+                return summary_df_single, plot_df, all_results
+
+            else: # Handles all other multi-run protocols
+                for i, run_spec in enumerate(protocol):
+                    label = run_spec["label"]
+                    current_probe_type = run_spec.get("probe_type", "seismic")
+                    dbg(f"--- Running Auto-Experiment: '{label}' ({i+1}/{len(protocol)}) ---")
+
+                    results = {}
+                    if current_probe_type == "act_titration":
+                        results = run_act_titration_probe(
+                            model_id=model_id, source_prompt_type=run_spec["source_prompt_type"],
+                            dest_prompt_type=run_spec["dest_prompt_type"], patch_steps=run_spec["patch_steps"],
+                            seed=seed, num_steps=num_steps, progress_callback=progress_callback,
+                        )
+                        summary_data.extend(results.get("titration_data", []))
+
+                    elif current_probe_type == "causal_surgery":
+                        results = run_causal_surgery_probe(
+                            model_id=model_id, source_prompt_type=run_spec["source_prompt_type"],
+                            dest_prompt_type=run_spec["dest_prompt_type"], patch_step=run_spec["patch_step"],
+                            seed=seed, num_steps=num_steps, progress_callback=progress_callback,
+                            reset_kv_cache_on_patch=run_spec.get("reset_kv_cache_on_patch", False)
+                        )
+                        stats = results.get("stats", {})
+                        patch_info = results.get("patch_info", {})
+                        summary_data.append({
+                            "Experiment": label, "Mean Delta": stats.get("mean_delta"),
+                            "Std Dev Delta": stats.get("std_delta"), "Max Delta": stats.get("max_delta"),
+                            "Introspective Report": results.get("introspective_report", "N/A"),
+                            "Patch Info": f"Source: {patch_info.get('source_prompt')}, Reset KV: {patch_info.get('kv_cache_reset')}"
+                        })
+
+                    elif current_probe_type == "triangulation":
+                        results = run_triangulation_probe(
+                            model_id=model_id, prompt_type=run_spec["prompt_type"], seed=seed, num_steps=num_steps,
+                            progress_callback=progress_callback, concept_to_inject=run_spec.get("concept", ""),
+                            injection_strength=run_spec.get("strength", 0.0),
+                        )
+                        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"),
+                            "Introspective Report": results.get("introspective_report", "N/A")
+                        })
+
+                    else: # seismic
+                        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.get("concept", ""), injection_strength=run_spec.get("strength", 0.0),
+                            progress_callback=progress_callback
+                        )
+                        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")
+                        })
+
+                    all_results[label] = results
+                    deltas = results.get("state_deltas", [])
+                    df = pd.DataFrame({"Step": range(len(deltas)), "Delta": deltas, "Experiment": label}) if deltas else pd.DataFrame()
+                    plot_data_frames.append(df)
+
+        summary_df = pd.DataFrame(summary_data)
 
-    return summary_df, plot_df, all_results
+        if probe_type == "act_titration":
+            plot_df = summary_df.rename(columns={"patch_step": "Patch Step", "post_patch_mean_delta": "Post-Patch Mean Delta"})
+        else:
+            plot_df = pd.concat(plot_data_frames, ignore_index=True) if plot_data_frames else pd.DataFrame()
+
+        if protocol and probe_type not in ["act_titration", "mechanistic_probe"]:
+            ordered_labels = [run['label'] for run in protocol]
+            if not summary_df.empty and 'Experiment' in summary_df.columns:
+                summary_df['Experiment'] = pd.Categorical(summary_df['Experiment'], categories=ordered_labels, ordered=True)
+                summary_df = summary_df.sort_values('Experiment')
+            if not plot_df.empty and 'Experiment' in plot_df.columns:
+                plot_df['Experiment'] = pd.Categorical(plot_df['Experiment'], categories=ordered_labels, ordered=True)
+                plot_df = plot_df.sort_values(['Experiment', 'Step'])
+
+        return summary_df, plot_df, all_results
+
+    finally:
+        if llm:
+            release_model(llm)

cognitive_mapping_probe/llm_iface.py

@@ -2,11 +2,12 @@ import os
 import torch
 import random
 import numpy as np
-from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
 from typing import Optional, List
 from dataclasses import dataclass, field
 
-from .utils import dbg
+# NEU: Importiere die zentrale cleanup-Funktion
+from .utils import dbg, cleanup_memory
 
 os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
@@ -17,34 +18,27 @@ class StableLLMConfig:
     layer_list: List[torch.nn.Module] = field(default_factory=list, repr=False)
 
 class LLM:
+    # __init__ und _populate_stable_config bleiben exakt wie in der vorherigen Version.
     def __init__(self, model_id: str, device: str = "auto", seed: int = 42):
         self.model_id = model_id
         self.seed = seed
         self.set_all_seeds(self.seed)
-
         token = os.environ.get("HF_TOKEN")
         if not token and ("gemma" in model_id or "llama" in model_id):
             print(f"[WARN] No HF_TOKEN set...", flush=True)
-
         kwargs = {"torch_dtype": torch.bfloat16} if torch.cuda.is_available() else {}
-
         dbg(f"Loading tokenizer for '{model_id}'...")
         self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True, token=token)
-
         dbg(f"Loading model '{model_id}' with kwargs: {kwargs}")
         self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device, token=token, **kwargs)
-
         try:
             self.model.set_attn_implementation('eager')
             dbg("Successfully set attention implementation to 'eager'.")
         except Exception as e:
             print(f"[WARN] Could not set 'eager' attention: {e}.", flush=True)
-
         self.model.eval()
         self.config = self.model.config
-
         self.stable_config = self._populate_stable_config()
-
         print(f"[INFO] Model '{model_id}' loaded on device: {self.model.device}", flush=True)
 
     def _populate_stable_config(self) -> StableLLMConfig:
@@ -53,7 +47,6 @@ class LLM:
             hidden_dim = self.model.get_input_embeddings().weight.shape[1]
         except AttributeError:
             hidden_dim = getattr(self.config, 'hidden_size', getattr(self.config, 'd_model', 0))
-
         num_layers = 0
         layer_list = []
         try:
@@ -63,26 +56,18 @@ class LLM:
                  layer_list = self.model.model.layers
             elif hasattr(self.model, 'transformer') and hasattr(self.model.transformer, 'h'):
                  layer_list = self.model.transformer.h
-
             if layer_list:
                 num_layers = len(layer_list)
         except (AttributeError, TypeError):
             pass
-
         if num_layers == 0:
             num_layers = getattr(self.config, 'num_hidden_layers', getattr(self.config, 'num_layers', 0))
-
         if hidden_dim <= 0 or num_layers <= 0 or not layer_list:
             dbg("--- CRITICAL: Failed to auto-determine model configuration. ---")
-            dbg(f"Detected hidden_dim: {hidden_dim}, num_layers: {num_layers}, found_layer_list: {bool(layer_list)}")
-            dbg("--- DUMPING MODEL ARCHITECTURE FOR DEBUGGING: ---")
             dbg(self.model)
-            dbg("--- END ARCHITECTURE DUMP ---")
-
         assert hidden_dim > 0, "Could not determine hidden dimension."
         assert num_layers > 0, "Could not determine number of layers."
         assert layer_list, "Could not find the list of transformer layers."
-
         dbg(f"Populated stable config: hidden_dim={hidden_dim}, num_layers={num_layers}")
         return StableLLMConfig(hidden_dim=hidden_dim, num_layers=num_layers, layer_list=layer_list)
 
@@ -97,30 +82,33 @@ class LLM:
         torch.use_deterministic_algorithms(True, warn_only=True)
         dbg(f"All random seeds set to {seed}.")
 
-    # --- NEU: Generische Text-Generierungs-Methode ---
     @torch.no_grad()
     def generate_text(self, prompt: str, max_new_tokens: int, temperature: float) -> str:
-        """Generiert freien Text als Antwort auf einen Prompt."""
-        self.set_all_seeds(self.seed) # Sorge für Reproduzierbarkeit
-
+        self.set_all_seeds(self.seed)
         messages = [{"role": "user", "content": prompt}]
         inputs = self.tokenizer.apply_chat_template(
             messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
         ).to(self.model.device)
-
         outputs = self.model.generate(
-            inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            do_sample=temperature > 0,
+            inputs, max_new_tokens=max_new_tokens, temperature=temperature, do_sample=temperature > 0,
         )
-
-        # Dekodiere nur die neu generierten Tokens
         response_tokens = outputs[0, inputs.shape[-1]:]
         return self.tokenizer.decode(response_tokens, skip_special_tokens=True)
 
 def get_or_load_model(model_id: str, seed: int) -> LLM:
+    """Lädt bei jedem Aufruf eine frische, isolierte Instanz des Modells."""
     dbg(f"--- Force-reloading model '{model_id}' for total run isolation ---")
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    cleanup_memory() # Bereinige Speicher, *bevor* ein neues Modell geladen wird.
     return LLM(model_id=model_id, seed=seed)
+
+# NEU: Explizite Funktion zum Freigeben von Ressourcen
+def release_model(llm: Optional[LLM]):
+    """
+    Gibt die Ressourcen eines LLM-Objekts explizit frei und ruft die zentrale
+    Speicherbereinigungs-Funktion auf.
+    """
+    if llm is None:
+        return
+    dbg(f"Releasing model instance for '{llm.model_id}'.")
+    del llm
+    cleanup_memory()

cognitive_mapping_probe/utils.py

@@ -1,15 +1,26 @@
 import os
 import sys
+import gc
+import torch
 
 # --- Centralized Debugging Control ---
-# To enable, set the environment variable: `export CMP_DEBUG=1`
 DEBUG_ENABLED = os.environ.get("CMP_DEBUG", "0") == "1"
 
 def dbg(*args, **kwargs):
-    """
-    A controlled debug print function. Only prints if DEBUG_ENABLED is True.
-    Ensures that debug output does not clutter production runs or HF Spaces logs
-    unless explicitly requested. Flushes output to ensure it appears in order.
-    """
+    """A controlled debug print function."""
     if DEBUG_ENABLED:
         print("[DEBUG]", *args, **kwargs, file=sys.stderr, flush=True)
+
+# --- NEU: Zentrale Funktion zur Speicherbereinigung ---
+def cleanup_memory():
+    """
+    Eine zentrale, global verfügbare Funktion zum Aufräumen von CPU- und GPU-Speicher.
+    Dies stellt sicher, dass die Speicherverwaltung konsistent und an einer einzigen Stelle erfolgt.
+    """
+    dbg("Cleaning up memory (centralized)...")
+    # Python's garbage collector
+    gc.collect()
+    # PyTorch's CUDA cache
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    dbg("Memory cleanup complete.")