Upload simple_inference.py

simple_inference.py

@@ -0,0 +1,96 @@
+import torch
+import os
+from modeling_physics_rl import PhysicsModel, Config
+
+def simple_inference():
+    print("🧪 Loading Physics Model (Controller + Adapters + WALT)...")
+    
+    # 1. Initialize Model Structure
+    model = PhysicsModel()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+
+    # 2. Load All Weights
+    try:
+        # A. Controller (The Brain)
+        if os.path.exists("final_physics_controller.pt"):
+            print("   loading controller...")
+            model.controller.load_state_dict(torch.load("final_physics_controller.pt", map_location=device))
+        
+        # B. Adapters (The Muscles)
+        if os.path.exists("final_flux_adapters.pt"):
+            print("   loading adapters...")
+            states = torch.load("final_flux_adapters.pt", map_location=device)
+            # Handle list vs dict
+            if isinstance(states, list):
+                for layer, state in zip(model.flux_layers, states): 
+                    layer.load_state_dict(state)
+            else:
+                model.flux_layers.load_state_dict(states)
+
+        # C. WALT Head (The Imagination)
+        if os.path.exists("final_walt_head.pt"):
+            print("   loading walt head...")
+            model.walt.load_state_dict(torch.load("final_walt_head.pt", map_location=device))
+            
+        print("✅ Model Assembled Successfully.")
+        
+    except Exception as e:
+        print(f"❌ Error loading weights: {e}")
+        return
+
+    # 3. Inference Loop
+    print("\n💡 Physics-Injected Inference (Type 'exit' to quit)")
+    print("   Input -> [Controller] -> Modulation -> [Adapters] -> Output\n")
+    
+    while True:
+        query = input("Query: ")
+        if query.lower() in ["exit", "quit"]:
+            break
+            
+        # Format for the trained distribution
+        prompt = f"User: {query}\nModel: "
+        inputs = model.tokenizer(prompt, return_tensors="pt").to(device)
+        
+        with torch.no_grad():
+            # A. Extract Physics Layout
+            h_init = model.get_embeddings(inputs.input_ids).to(Config.DTYPE)
+            
+            # B. Controller Decision
+            modulation = model.controller(h_init)
+            mod_norm = torch.norm(modulation).item()
+            
+            # C. Inject Physics (Activate Adapters)
+            model.set_active_modulation(modulation)
+            
+            # D. (Optional) WALT Prediction (Just to verify it runs)
+            # z, z_next = model.walt(h_init) 
+            
+            # E. Generate
+            out_ids = model.llm.generate(
+                **inputs,
+                max_new_tokens=128,
+                do_sample=True,
+                temperature=0.6,
+                top_p=0.9,
+                repetition_penalty=1.2,
+                pad_token_id=model.tokenizer.eos_token_id
+            )
+            
+            # Reset
+            model.clear_modulation()
+            
+        response = model.tokenizer.decode(out_ids[0], skip_special_tokens=True)
+        
+        # Clean Prompt
+        if response.startswith(prompt):
+            response = response[len(prompt):].strip()
+        elif "Model:" in response:
+            response = response.split("Model:")[-1].strip()
+            
+        print(f"Response: {response}")
+        print(f"   [Physics Injection Intensity: {mod_norm:.4f}]\n")
+
+if __name__ == "__main__":
+    simple_inference()