Upload 2 files

app.py

@@ -1,55 +1,70 @@
 import gradio as gr
 import torch
-import tiktoken
-from train_shakespeare import GPT, GPTConfig, generate, get_autocast_device
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-# Initialize model and tokenizer
-def init_model():
-    model = GPT(GPTConfig())
-    checkpoint = torch.load('model/best_model.pt', map_location='cpu')
-    model.load_state_dict(checkpoint['model_state_dict'])
+# Load model and tokenizer from Hugging Face
+def load_model():
+    model_name = "aayushraina/gpt2shakespeare"
+    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+    model = GPT2LMHeadModel.from_pretrained(model_name)
     model.eval()
-    return model
+    return model, tokenizer
 
-enc = tiktoken.get_encoding("gpt2")
-model = init_model()
-
-def generate_text(prompt, max_length=500, temperature=0.8, top_k=40):
-    # Tokenize input
-    input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)
+# Text generation function
+def generate_text(prompt, max_length=500, temperature=0.8, top_k=40, top_p=0.9):
+    # Encode the input prompt
+    input_ids = tokenizer.encode(prompt, return_tensors='pt')
     
     # Generate text
     with torch.no_grad():
-        output_ids = generate(
-            model=model,
-            idx=input_ids,
-            max_new_tokens=max_length,
+        output = model.generate(
+            input_ids,
+            max_length=max_length,
             temperature=temperature,
             top_k=top_k,
-            device='cpu'  # Force CPU for Spaces
+            top_p=top_p,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            num_return_sequences=1
         )
     
-    # Decode and return generated text
-    return enc.decode(output_ids[0].tolist())
+    # Decode and return the generated text
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return generated_text
+
+# Load model and tokenizer globally
+print("Loading model and tokenizer...")
+model, tokenizer = load_model()
+print("Model loaded successfully!")
 
 # Create Gradio interface
 demo = gr.Interface(
     fn=generate_text,
     inputs=[
-        gr.Textbox(label="Enter your prompt", placeholder="Start your text here..."),
+        gr.Textbox(label="Enter your prompt", placeholder="Start your text here...", lines=2),
         gr.Slider(minimum=10, maximum=1000, value=500, step=10, label="Maximum Length"),
         gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Temperature"),
-        gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-k")
+        gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-k"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p"),
     ],
-    outputs=gr.Textbox(label="Generated Text"),
+    outputs=gr.Textbox(label="Generated Text", lines=10),
     title="Shakespeare-style Text Generator",
-    description="Generate Shakespeare-style text using a fine-tuned GPT-2 model",
+    description="""Generate Shakespeare-style text using a fine-tuned GPT-2 model.
+    
+    Parameters:
+    - Temperature: Higher values make the output more random, lower values more focused
+    - Top-k: Number of highest probability vocabulary tokens to keep for top-k filtering
+    - Top-p: Cumulative probability for nucleus sampling
+    """,
     examples=[
-        ["First Citizen:", 500, 0.8, 40],
-        ["To be, or not to be,", 500, 0.8, 40],
-        ["Friends, Romans, countrymen,", 500, 0.8, 40]
+        ["First Citizen:", 500, 0.8, 40, 0.9],
+        ["To be, or not to be,", 500, 0.8, 40, 0.9],
+        ["Friends, Romans, countrymen,", 500, 0.8, 40, 0.9],
+        ["O Romeo, Romeo,", 500, 0.8, 40, 0.9],
+        ["Now is the winter of our discontent", 500, 0.8, 40, 0.9]
     ]
 )
 
+# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch() 

train_shakespeare.py

@@ -147,92 +147,97 @@ class DataLoaderLite:
             self.current_position = 0
         return x, y
 
-# Device configuration
-device = 'cpu'
-if torch.cuda.is_available():
-    device = 'cuda'
-elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-    device = "mps"
-print(f"using device: {device}")
-
-# Set random seed
-torch.manual_seed(1337)
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(1337)
-
-# Initialize model and move to device
-model = GPT(GPTConfig())
-model.to(device)
-
-# Initialize data loader
-train_loader = DataLoaderLite(B=4, T=32)
-
-# Training settings
-learning_rate = 3e-4
-num_iters = 100000  # Increased to 100000
-eval_interval = 50   # Evaluate every 50 iterations
-best_loss = float('inf')
-checkpoint_dir = 'checkpoints'
-os.makedirs(checkpoint_dir, exist_ok=True)
-
-# Initialize optimizer
-optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
-
-print(f"\n=== Starting Training ===")
-print(f"Total iterations: {num_iters}")
-print(f"Evaluation interval: {eval_interval}")
-print(f"Learning rate: {learning_rate}")
-
-# Training loop
-for iter in range(num_iters):
-    # Get batch
-    x, y = train_loader.next_batch()
-    x, y = x.to(device), y.to(device)
-    
-    # Forward pass
-    optimizer.zero_grad()
-    logits, loss = model(x, y)
-    
-    # Backward pass
-    loss.backward()
-    optimizer.step()
-    
-    # Log progress every 50 iterations
-    if iter % eval_interval == 0:
-        current_loss = loss.item()
-        print(f'step {iter}, loss: {current_loss:.4f}')
+
+# write the main block
+
+if __name__ == "__main__":
+
+    # Device configuration
+    device = 'cpu'
+    if torch.cuda.is_available():
+        device = 'cuda'
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    print(f"using device: {device}")
+
+    # Set random seed
+    torch.manual_seed(1337)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(1337)
+
+    # Initialize model and move to device
+    model = GPT(GPTConfig())
+    model.to(device)
+
+    # Initialize data loader
+    train_loader = DataLoaderLite(B=4, T=32)
+
+    # Training settings
+    learning_rate = 3e-4
+    num_iters = 100000  # Increased to 100000
+    eval_interval = 50   # Evaluate every 50 iterations
+    best_loss = float('inf')
+    checkpoint_dir = 'checkpoints'
+    os.makedirs(checkpoint_dir, exist_ok=True)
+
+    # Initialize optimizer
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+    print(f"\n=== Starting Training ===")
+    print(f"Total iterations: {num_iters}")
+    print(f"Evaluation interval: {eval_interval}")
+    print(f"Learning rate: {learning_rate}")
+
+    # Training loop
+    for iter in range(num_iters):
+        # Get batch
+        x, y = train_loader.next_batch()
+        x, y = x.to(device), y.to(device)
+        
+        # Forward pass
+        optimizer.zero_grad()
+        logits, loss = model(x, y)
+        
+        # Backward pass
+        loss.backward()
+        optimizer.step()
         
-        # Save if this is the best model so far
-        if current_loss < best_loss:
-            best_loss = current_loss
-            checkpoint_path = os.path.join(checkpoint_dir, f'model_step_{iter}_loss_{current_loss:.4f}.pt')
-            torch.save({
-                'iter': iter,
-                'model_state_dict': model.state_dict(),
-                'optimizer_state_dict': optimizer.state_dict(),
-                'loss': current_loss,
-                'best_loss': best_loss,
-            }, checkpoint_path)
-            print(f'New best model saved! Loss: {current_loss:.4f}')
+        # Log progress every 50 iterations
+        if iter % eval_interval == 0:
+            current_loss = loss.item()
+            print(f'step {iter}, loss: {current_loss:.4f}')
             
-            # Also save as best model
-            torch.save({
-                'iter': iter,
-                'model_state_dict': model.state_dict(),
-                'optimizer_state_dict': optimizer.state_dict(),
-                'loss': current_loss,
-                'best_loss': best_loss,
-            }, 'best_model.pt')
-
-print("\n=== Training Complete ===")
-print(f"Best loss achieved: {best_loss:.4f}")
-
-# Save final model
-final_path = os.path.join(checkpoint_dir, 'model_final.pt')
-torch.save({
-    'iter': num_iters-1,
-    'model_state_dict': model.state_dict(),
-    'optimizer_state_dict': optimizer.state_dict(),
-    'loss': loss.item(),
-    'best_loss': best_loss,
-}, final_path)
+            # Save if this is the best model so far
+            if current_loss < best_loss:
+                best_loss = current_loss
+                checkpoint_path = os.path.join(checkpoint_dir, f'model_step_{iter}_loss_{current_loss:.4f}.pt')
+                torch.save({
+                    'iter': iter,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'loss': current_loss,
+                    'best_loss': best_loss,
+                }, checkpoint_path)
+                print(f'New best model saved! Loss: {current_loss:.4f}')
+                
+                # Also save as best model
+                torch.save({
+                    'iter': iter,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'loss': current_loss,
+                    'best_loss': best_loss,
+                }, 'best_model.pt')
+
+    print("\n=== Training Complete ===")
+    print(f"Best loss achieved: {best_loss:.4f}")
+
+    # Save final model
+    final_path = os.path.join(checkpoint_dir, 'model_final.pt')
+    torch.save({
+        'iter': num_iters-1,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'loss': loss.item(),
+        'best_loss': best_loss,
+    }, final_path)