Upload app.py

app.py

@@ -114,35 +114,97 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 config = GPTConfig()
 model = GPT(config)
 
+model_loaded = False
+
 # Try to load model from HuggingFace Model Hub first, then local file
 try:
     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
+    repo_id = os.getenv('HF_MODEL_REPO', 'shwethd/gpt2-shakespeare-124m')
     
     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}")
+        print(f"Attempting to load from HuggingFace Hub: {repo_id}")
+        
+        # Try SafeTensors first (more secure, no pickle issues)
+        try:
+            from safetensors.torch import load_file
+            try:
+                model_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename="model.safetensors",
+                    cache_dir=None
+                )
+                state_dict = load_file(model_path, device=device)
+                model.load_state_dict(state_dict)
+                model_loaded = True
+                print(f"✅ Model loaded successfully from SafeTensors: {repo_id}")
+            except Exception as e:
+                print(f"SafeTensors not found ({e}), trying .pt file...")
+                # Fallback to .pt file
+                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)
+                
+                # Handle different checkpoint formats
+                if 'model_state_dict' in checkpoint:
+                    model.load_state_dict(checkpoint['model_state_dict'])
+                elif 'state_dict' in checkpoint:
+                    model.load_state_dict(checkpoint['state_dict'])
+                else:
+                    # If checkpoint is the state dict itself
+                    model.load_state_dict(checkpoint)
+                
+                model_loaded = True
+                print(f"✅ Model loaded successfully from HuggingFace Hub: {repo_id}")
+        except ImportError:
+            # safetensors not installed, use .pt file
+            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)
+            
+            # Handle different checkpoint formats
+            if 'model_state_dict' in checkpoint:
+                model.load_state_dict(checkpoint['model_state_dict'])
+            elif 'state_dict' in checkpoint:
+                model.load_state_dict(checkpoint['state_dict'])
+            else:
+                # If checkpoint is the state dict itself
+                model.load_state_dict(checkpoint)
+            
+            model_loaded = True
+            print(f"✅ Model loaded successfully 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")
+        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)
+            if 'model_state_dict' in checkpoint:
+                model.load_state_dict(checkpoint['model_state_dict'])
+            elif 'state_dict' in checkpoint:
+                model.load_state_dict(checkpoint['state_dict'])
+            else:
+                model.load_state_dict(checkpoint)
+            model_loaded = True
+            print("✅ Model loaded from local checkpoint")
+        except Exception as e2:
+            print(f"❌ Could not load from local file either: {e2}")
 except FileNotFoundError:
-    print("Warning: Model checkpoint not found. Using untrained model.")
-    # Model will be randomly initialized - not ideal but won't crash
+    print("❌ Warning: Model checkpoint not found. Using untrained model.")
 except Exception as e:
-    print(f"Error loading model: {e}")
-    print("Using untrained model as fallback.")
+    print(f"❌ Error loading model: {e}")
+    print("⚠️ Using untrained model as fallback - output will be random!")
+
+if not model_loaded:
+    print("⚠️ WARNING: Model is using random weights! Generation will be nonsensical.")
+    print("Please ensure model_checkpoint_final.pt is uploaded to HuggingFace Model Hub.")
 
 model.to(device)
 model.eval()
@@ -154,21 +216,47 @@ enc = tiktoken.get_encoding('gpt2')
 def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
     """Generate text from prompt"""
     try:
+        if not model_loaded:
+            return "❌ Error: Model not loaded correctly. Please check that model_checkpoint_final.pt is uploaded to HuggingFace Model Hub (shwethd/gpt2-shakespeare-124m)."
+        
+        # Validate inputs
+        if not prompt or len(prompt.strip()) == 0:
+            return "Please enter a prompt."
+        
+        temperature = max(0.1, min(2.0, temperature))  # Clamp temperature
+        top_k = max(1, min(100, int(top_k)))  # Clamp top_k
+        max_new_tokens = max(1, min(200, int(max_new_tokens)))  # Clamp max tokens
+        
         # Encode prompt
         tokens = enc.encode(prompt)
+        if len(tokens) == 0:
+            return "Error: Could not encode prompt."
+        
         tokens = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)
         
         # Generate
         with torch.no_grad():
-            for _ in range(max_new_tokens):
+            for i in range(max_new_tokens):
                 # Forward pass
                 logits, _ = model(tokens)
-                logits = logits[:, -1, :] / temperature
+                logits = logits[:, -1, :] / max(temperature, 0.1)  # Avoid division by zero
+                
+                # Apply top-k filtering
+                if top_k < logits.size(-1):
+                    topk_logits, topk_indices = torch.topk(logits, top_k, dim=-1)
+                    # Create filtered logits
+                    filtered_logits = torch.full_like(logits, float('-inf'))
+                    filtered_logits.scatter_(-1, topk_indices, topk_logits)
+                    logits = filtered_logits
+                
+                # Sample from distribution
+                probs = F.softmax(logits, dim=-1)
+                
+                # Avoid NaN
+                if torch.isnan(probs).any():
+                    probs = torch.ones_like(probs) / probs.size(-1)
                 
-                # Top-k sampling
-                topk_probs, topk_indices = torch.topk(F.softmax(logits, dim=-1), top_k, dim=-1)
-                ix = torch.multinomial(topk_probs, 1)
-                next_token = torch.gather(topk_indices, -1, ix)
+                next_token = torch.multinomial(probs, 1)
                 
                 # Append to sequence
                 tokens = torch.cat([tokens, next_token], dim=1)
@@ -181,14 +269,21 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
         generated_text = enc.decode(tokens[0].tolist())
         return generated_text
     except Exception as e:
-        return f"Error: {str(e)}"
+        import traceback
+        return f"❌ Error during generation: {str(e)}\n\nPlease check:\n1. Model is uploaded to HuggingFace Model Hub\n2. Repository name is correct: shwethd/gpt2-shakespeare-124m\n3. File name is exactly: model_checkpoint_final.pt"
 
 
 # Create Gradio interface
 with gr.Blocks(title="GPT-2 124M Shakespeare Model") as demo:
-    gr.Markdown("""
+    # Status indicator
+    status_color = "🟢" if model_loaded else "🔴"
+    status_text = "Model loaded successfully!" if model_loaded else "⚠️ Model not loaded - check HuggingFace Model Hub!"
+    
+    gr.Markdown(f"""
     # 🎭 GPT-2 124M Shakespeare Language Model
     
+    {status_color} **Status:** {status_text}
+    
     This is a 124M parameter decoder-only transformer model trained on Shakespeare's complete works.
     
     **Training Results:**
@@ -197,6 +292,8 @@ with gr.Blocks(title="GPT-2 124M Shakespeare Model") as demo:
     - Training Steps: 1,637
     
     Enter a prompt below to generate Shakespeare-style text!
+    
+    {"⚠️ **Note:** If you see garbled/random text, the model may not have loaded correctly. Check the logs and ensure the model is uploaded to HuggingFace Model Hub: `shwethd/gpt2-shakespeare-124m`" if not model_loaded else ""}
     """)
     
     with gr.Row():
@@ -238,7 +335,7 @@ with gr.Blocks(title="GPT-2 124M Shakespeare Model") as demo:
             )
     
     # Example prompts
-    gr.Markdown("### Example Prompts:")
+    gr.Markdown("### Example Prompts (Click to try):")
     examples = gr.Examples(
         examples=[
             ["First Citizen:"],
@@ -246,6 +343,15 @@ with gr.Blocks(title="GPT-2 124M Shakespeare Model") as demo:
             ["To be or not"],
             ["HAMLET:"],
             ["MACBETH:"],
+            ["JULIET:"],
+            ["KING:"],
+            ["LADY MACBETH:"],
+            ["OTHELLO:"],
+            ["What light through yonder"],
+            ["All the world's a stage"],
+            ["Double, double toil and trouble"],
+            ["Friends, Romans, countrymen"],
+            ["A rose by any other name"],
         ],
         inputs=prompt_input
     )