Upload app.py

app.py

@@ -138,6 +138,9 @@ try:
                 )
                 state_dict = load_file(model_path, device=device)
                 model.load_state_dict(state_dict)
+                # Restore weight sharing (broken during SafeTensors conversion)
+                # lm_head.weight and transformer.wte.weight should share memory
+                model.transformer.wte.weight = model.lm_head.weight
                 model_loaded = True
                 print(f"✅ Model loaded successfully from SafeTensors: {repo_id}")
             except Exception as e:
@@ -148,7 +151,9 @@ try:
                     filename="model_checkpoint_final.pt",
                     cache_dir=None
                 )
-                checkpoint = torch.load(model_path, map_location=device)
+                # PyTorch 2.6+ requires weights_only=False for custom classes
+                # This is safe since we trust our own trained model
+                checkpoint = torch.load(model_path, map_location=device, weights_only=False)
                 
                 # Handle different checkpoint formats
                 if 'model_state_dict' in checkpoint:
@@ -168,7 +173,8 @@ try:
                 filename="model_checkpoint_final.pt",
                 cache_dir=None
             )
-            checkpoint = torch.load(model_path, map_location=device)
+            # PyTorch 2.6+ requires weights_only=False for custom classes
+            checkpoint = torch.load(model_path, map_location=device, weights_only=False)
             
             # Handle different checkpoint formats
             if 'model_state_dict' in checkpoint:
@@ -185,7 +191,8 @@ try:
         print(f"⚠️ Could not load from Hub ({e}), trying local file...")
         try:
             # Fallback to local file
-            checkpoint = torch.load('model_checkpoint_final.pt', map_location=device)
+            # PyTorch 2.6+ requires weights_only=False for custom classes
+            checkpoint = torch.load('model_checkpoint_final.pt', map_location=device, weights_only=False)
             if 'model_state_dict' in checkpoint:
                 model.load_state_dict(checkpoint['model_state_dict'])
             elif 'state_dict' in checkpoint:
@@ -369,5 +376,6 @@ with gr.Blocks(title="GPT-2 124M Shakespeare Model") as demo:
     """)
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    # Don't use share=True on HuggingFace Spaces
+    demo.launch()