Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 import yaml
 import json
 from tokenizers import Tokenizer
+from collections import OrderedDict
 
 # --- 1. Load Custom Model Code ---
 # This import now works because we have the correct models/hrm/ structure
@@ -26,8 +27,24 @@ model_config.update({
     'vocab_size': tokenizer.get_vocab_size()
 })
 model = HierarchicalReasoningModel_ACTV1(config_dict=model_config)
-model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
-model.eval()
+
+# --- MODIFICATION: Clean the state dict keys before loading ---
+# Load the original state dict
+original_state_dict = torch.load('pytorch_model.bin', map_location='cpu')
+# Create a new state dict with the corrected keys
+new_state_dict = OrderedDict()
+for k, v in original_state_dict.items():
+    if k.startswith('_orig_mod.model.'):
+        name = k[len('_orig_mod.model.'):] # remove the prefix
+        new_state_dict[name] = v
+    else:
+        new_state_dict[k] = v
+
+# Load the cleaned state dict
+model.load_state_dict(new_state_dict)
+# --- END MODIFICATION ---
+
+model.eval() # Set the model to evaluation mode
 print("Model loaded successfully!")
 
 # --- 4. Define the Inference Function ---