Upload benchmark_physics.py with huggingface_hub

benchmark_physics.py

@@ -0,0 +1,222 @@
+
+import torch
+import logging
+import os
+import glob
+from config_physics import Config
+from modeling_physics_rl import PhysicsModel
+
+# Setup logging
+logging.basicConfig(level=logging.ERROR)
+
+def load_models():
+    """
+    Loads two versions of the model:
+    1. Flux Model: With trained Controller & Adapters active.
+    2. Base Model: The exact same model but with modulation forced to ZERO.
+    """
+    print("⏳ Loading Physics Model...")
+    model = PhysicsModel()
+    
+    # Move to GPU if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    print(f"   Using Device: {device}")
+    
+    # Load Weights
+    # Define search paths
+    search_paths = [
+        ".", 
+        "/kaggle/input/worldmodels/physics_model",
+        "/kaggle/working/physics_model"
+    ]
+    
+    # Check for weights
+    final_path = None
+    for p in search_paths:
+        fpath = os.path.join(p, "final_physics_controller.pt")
+        if os.path.exists(fpath):
+            final_path = p
+            break
+            
+    try:
+        if final_path:
+            print(f"   Loading Final Weights from {final_path}...")
+            model.controller.load_state_dict(torch.load(os.path.join(final_path, "final_physics_controller.pt"), map_location=model.llm.device))
+            model.walt.load_state_dict(torch.load(os.path.join(final_path, "final_walt_head.pt"), map_location=model.llm.device))
+            
+            # Load Adapters
+            adapter_path = os.path.join(final_path, "final_flux_adapters.pt")
+            if os.path.exists(adapter_path):
+                print("   Loading Flux Adapters...")
+                adapter_states = torch.load(adapter_path, map_location=model.llm.device)
+                for layer, state in zip(model.flux_layers, adapter_states):
+                    layer.load_state_dict(state)
+            else:
+                print("   ⚠️ Startled: Final adapters not found! Modulation might be dead.")
+
+        else:
+            # Fallback to latest checkpoint
+            checkpoints = []
+            for p in search_paths:
+                checkpoints.extend(glob.glob(os.path.join(p, "checkpoint_epoch_*.pt")))
+            
+            if checkpoints:
+                latest_ckpt = max(checkpoints, key=os.path.getctime)
+                print(f"   ⚠️ 'final' weights not found. Loading latest checkpoint: {latest_ckpt}")
+                ckpt_data = torch.load(latest_ckpt, map_location=model.llm.device)
+                
+                # Check point uses specific keys, not full model_state_dict
+                if 'controller_state_dict' in ckpt_data:
+                    model.controller.load_state_dict(ckpt_data['controller_state_dict'])
+                    model.walt.load_state_dict(ckpt_data['walt_state_dict'])
+                    
+                    if 'adapters_state_dict' in ckpt_data:
+                        print("   Loading Flux Adapters from Checkpoint...")
+                        for layer, state in zip(model.flux_layers, ckpt_data['adapters_state_dict']):
+                            layer.load_state_dict(state)
+                else:
+                    # Fallback if we accidentally saved it differently in a previous run
+                    model.load_state_dict(ckpt_data['model_state_dict'], strict=False)
+            else:
+                raise FileNotFoundError("No 'final_physics_controller.pt' or 'checkpoint_epoch_*.pt' found.")
+            
+        print("✅ Weights Loaded.")
+    except Exception as e:
+        print(f"⚠️ Warning: Could not load weights: {e}")
+
+    model.eval()
+    return model
+
+def run_benchmark():
+    model = load_models()
+    
+    # Health Check
+    try:
+        if hasattr(model.flux_layers[0], 'lora_B'):
+            lb_norm = model.flux_layers[0].lora_B.norm().item()
+            print(f"\n🔍 Health Check - First Adapter LoRA_B Norm: {lb_norm:.6f}")
+            if lb_norm == 0:
+                print("   ❌ WARNING: LoRA weights are ZERO. Training failed to update weights.")
+            else:
+                print("   ✅ Weights are LEARNED (Non-Zero).")
+    except: pass
+
+    test_cases = [
+        # --- TYPE A: QUALITATIVE (Concept Checks) ---
+        "I release a heavy steel marble from a height of one meter in a zero-gravity environment.",
+        "I drop a plastic camping plate onto a marble floor from waist height.",
+        "I shine a red laser beam through a glass prism.", 
+        
+        # --- TYPE B: QUANTITATIVE (Math & Engineering) ---
+        "A 2kg block slides down a frictionless ramp of height 5m. Calculate its velocity at the bottom. (g=9.8 m/s^2)",
+        "A car accelerates from 0 to 20 m/s in 4 seconds. What is the average acceleration?",
+        "A one-meter-long flexible cable lies at rest on a frictionless table, with 5 cm hanging over the edge. At what time will the cable completely slide off the table?",
+        "If I mix 100g of ice at 0°C with 100g of water at 80°C, what is the final temperature? (Specific heat of water = 4.18 J/g°C)",
+    ]
+    
+    results = []
+    
+    print("\n" + "="*50)
+    print(" 🧪 Physics Benchmark: Base vs Flux")
+    print("="*50)
+    
+    for prompt in test_cases:
+        full_prompt = f"User: {prompt}\nModel:"
+        inputs = model.tokenizer(full_prompt, return_tensors="pt").to(model.llm.device)
+        
+        # --- Run 1: Base Model (No Modulation) ---
+        model.clear_modulation() # Ensure no modulation
+        # We can simulate "Base" by simply NOT calling set_active_modulation
+        # Or by setting modulation to all zeros.
+        # Let's set to zeros to be explicit.
+        zero_mod = torch.zeros(1, Config.MODULATION_DIM).to(model.llm.device).to(Config.DTYPE)
+        model.set_active_modulation(zero_mod)
+        
+        out_base = model.llm.generate(**inputs, max_new_tokens=100, max_length=Config.MAX_LENGTH, do_sample=False) # Greedy for base
+        text_base = model.tokenizer.decode(out_base[0], skip_special_tokens=True).replace(full_prompt, "").strip()
+        
+        # --- Run 2: Flux Model (With RL Modulation) ---
+        model.clear_modulation()
+        
+        # Thinking Step
+        with torch.no_grad():
+            h_init = model.get_embeddings(inputs.input_ids).to(Config.DTYPE)
+            modulation = model.controller(h_init)
+            
+            # Analyze Modulation strength
+            mod_mag = modulation.norm().item()
+            
+            model.set_active_modulation(modulation)
+            
+            # --- Debug Trace (First 3 tokens) ---
+            try:
+                print("\n   🔍 Generation Trace (First 3 Steps):")
+                trace_input = inputs.input_ids.clone()
+                for i in range(3):
+                    # Base (No Mod)
+                    model.clear_modulation()
+                    out_base = model.llm.model(trace_input)
+                    base_norm = out_base.last_hidden_state[:,-1,:].norm().item()
+                    
+                    # Flux (Modulated)
+                    model.set_active_modulation(modulation)
+                    out_liq = model.llm.model(trace_input)
+                    liq_norm = out_liq.last_hidden_state[:,-1,:].norm().item()
+                    
+                    # Difference
+                    diff = out_liq.last_hidden_state[:,-1,:] - out_base.last_hidden_state[:,-1,:]
+                    diff_norm = diff.norm().item()
+                    ratio = (diff_norm / base_norm) * 100
+                    
+                    print(f"      Step {i}: Base={base_norm:.2f} | Flux={liq_norm:.2f} | Diff={diff_norm:.4f} ({ratio:.2f}%)")
+                    
+                    # Advance one step (Greedy)
+                    # Use internal lm_head to get logits
+                    logits = model.llm.lm_head(out_liq.last_hidden_state[:,-1,:].unsqueeze(0))
+                    # Check dim
+                    if logits.dim() == 3: logits = logits[:, -1, :]
+                    
+                    next_token = torch.argmax(logits, dim=-1).unsqueeze(0)
+                    token_str = model.tokenizer.decode(next_token[0])
+                    print(f"      Selected Token: '{token_str}'")
+
+                    if next_token.dim() == 1: next_token = next_token.unsqueeze(0)
+                        
+                    trace_input = torch.cat([trace_input, next_token], dim=1)
+            except Exception as e:
+                print(f"   ⚠️ Debug Trace Failed: {e}")
+            
+            # Reset for actual generation
+            model.clear_modulation()
+            model.set_active_modulation(modulation)
+            
+        out_liquid = model.llm.generate(**inputs, max_new_tokens=100, max_length=Config.MAX_LENGTH, do_sample=True, temperature=0.01)
+        text_liquid = model.tokenizer.decode(out_liquid[0], skip_special_tokens=True).replace(full_prompt, "").strip()
+        
+        # Store Result
+        res = {
+            "Prompt": prompt,
+            "Base": text_base,
+            "Flux": text_liquid,
+            "Modulation_Norm": mod_mag
+        }
+        results.append(res)
+        
+        print(f"\n📝 {prompt}")
+        print(f"   🧊 Base:   {text_base[:100]}...")
+        print(f"   💧 Flux:   {text_liquid[:100]}... (Mod Norm: {mod_mag:.2f})")
+        
+    # Save detailed report
+    with open("benchmark_results.txt", "w") as f:
+        for r in results:
+            f.write(f"Prompt: {r['Prompt']}\n")
+            f.write(f"Base Model: {r['Base']}\n")
+            f.write(f"Flux Model: {r['Flux']}\n")
+            f.write(f"Modulation Strength: {r['Modulation_Norm']:.4f}\n")
+            f.write("-" * 30 + "\n")
+            
+    print("\n✅ Benchmark Complete. Saved to benchmark_results.txt")
+
+if __name__ == "__main__":
+    run_benchmark()