Upload app.py

app.py

@@ -114,14 +114,35 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 config = GPTConfig()
 model = GPT(config)
 
-# Try to load model (works both locally and on HuggingFace)
+# Try to load model from HuggingFace Model Hub first, then local file
 try:
-    checkpoint = torch.load('model_checkpoint_final.pt', map_location=device)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    print("Model loaded from checkpoint")
+    from huggingface_hub import hf_hub_download
+    import os
+    
+    # Try to get model path from environment variable or use default
+    repo_id = os.getenv('HF_MODEL_REPO', 'YOUR_USERNAME/gpt2-shakespeare-124m')  # Update with your repo
+    
+    try:
+        model_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="model_checkpoint_final.pt",
+            cache_dir=None
+        )
+        checkpoint = torch.load(model_path, map_location=device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Model loaded from HuggingFace Hub: {repo_id}")
+    except Exception as e:
+        print(f"Could not load from Hub ({e}), trying local file...")
+        # Fallback to local file
+        checkpoint = torch.load('model_checkpoint_final.pt', map_location=device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print("Model loaded from local checkpoint")
 except FileNotFoundError:
     print("Warning: Model checkpoint not found. Using untrained model.")
     # Model will be randomly initialized - not ideal but won't crash
+except Exception as e:
+    print(f"Error loading model: {e}")
+    print("Using untrained model as fallback.")
 
 model.to(device)
 model.eval()