Update app.py

app.py

@@ -1,12 +1,10 @@
-import os
-import math
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
-import gradio as gr
 import tiktoken
+import gradio as gr
 
-# GPT model code
+# Define the model architecture
 class GPTConfig:
     def __init__(self):
         self.block_size = 1024
@@ -100,7 +98,7 @@ class GPT(nn.Module):
         
         return logits, loss
 
-# Updated load_model function
+# Load the model
 def load_model(model_path):
     config = GPTConfig()
     model = GPT(config)
@@ -110,46 +108,63 @@ def load_model(model_path):
     print("Checkpoint keys:", checkpoint.keys())  # Debug print
     
     if 'model_state_dict' in checkpoint:
-        # If the checkpoint contains a 'model_state_dict' key, use that
         model.load_state_dict(checkpoint['model_state_dict'])
     else:
-        # Otherwise, try to load the state dict directly
         model.load_state_dict(checkpoint)
     
     model.eval()
     return model
 
-# Load the trained model
+# Load the model
 model = load_model('gpt_5000.pt')  # Replace with the actual path to your .pt file
 enc = tiktoken.get_encoding('gpt2')
 
-def generate_text(prompt, max_length=100, temperature=0.7):
+# Improved text generation function
+def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
     input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)
+    generated = []
     
     with torch.no_grad():
         for _ in range(max_length):
             outputs, _ = model(input_ids)
-            next_token_logits = outputs[:, -1, :] / temperature
-            next_token = torch.multinomial(torch.softmax(next_token_logits, dim=-1), num_samples=1)
+            next_token_logits = outputs[:, -1, :]
+            
+            # Apply temperature
+            next_token_logits = next_token_logits / temperature
+            
+            # Apply top-k filtering
+            top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k, dim=-1)
+            next_token_probs = F.softmax(top_k_logits, dim=-1)
+            
+            # Sample from the filtered distribution
+            next_token_index = torch.multinomial(next_token_probs, num_samples=1)
+            next_token = top_k_indices.gather(-1, next_token_index)
+            
             input_ids = torch.cat([input_ids, next_token], dim=-1)
+            generated.append(next_token.item())
             
-            if next_token.item() == enc.encode('\n')[0]:
+            # Stop if we generate a newline, but only after generating at least 20 tokens
+            if next_token.item() == enc.encode('\n')[0] and len(generated) > 20:
                 break
     
-    generated_text = enc.decode(input_ids[0].tolist())
-    return generated_text
+    generated_text = enc.decode(generated)
+    return prompt + generated_text
 
 # Gradio interface
+def gradio_generate(prompt, max_length, temperature, top_k):
+    return generate_text(prompt, max_length, temperature, top_k)
+
 iface = gr.Interface(
-    fn=generate_text,
+    fn=gradio_generate,
     inputs=[
         gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
-        gr.Slider(minimum=10, maximum=200, value=100, step=1, label="Max Length"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
+        gr.Slider(minimum=20, maximum=500, value=100, step=1, label="Max Length"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k")
     ],
     outputs=gr.Textbox(label="Generated Text"),
-    title="GPT-2 Text Generator",
-    description="Enter a prompt and generate text using a fine-tuned GPT-2 model."
+    title="GPT Text Generator",
+    description="Enter a prompt and adjust parameters to generate text using a fine-tuned GPT model."
 )
 
 # Launch the app