Update app.py

app.py

@@ -120,6 +120,15 @@ model = load_model('gpt_model.pth')  # Replace with the actual path to your .pt
 enc = tiktoken.get_encoding('gpt2')
 
 # Improved text generation function
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import tiktoken
+import gradio as gr
+
+# [Your existing model code remains unchanged]
+
+# Modified text generation function to yield tokens
 def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
     input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)
     generated = []
@@ -128,32 +137,70 @@ def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
         for _ in range(max_length):
             outputs, _ = model(input_ids)
             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())
             
-            # Stop if we generate a newline, but only after generating at least 20 tokens
+            yield enc.decode([next_token.item()])
+            
             if next_token.item() == enc.encode('\n')[0] and len(generated) > 20:
                 break
-    
-    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)
 
+# Custom CSS for the animation effect
+custom_css = """
+<style>
+    .output-box {
+        border: 1px solid #e0e0e0;
+        border-radius: 8px;
+        padding: 20px;
+        font-family: Arial, sans-serif;
+        line-height: 1.6;
+        height: 300px;
+        overflow-y: auto;
+        background-color: #f9f9f9;
+    }
+    .blinking-cursor {
+        display: inline-block;
+        width: 10px;
+        height: 20px;
+        background-color: #333;
+        animation: blink 0.7s infinite;
+    }
+    @keyframes blink {
+        0% { opacity: 0; }
+        50% { opacity: 1; }
+        100% { opacity: 0; }
+    }
+</style>
+"""
+
+# JavaScript for the typing animation
+js_code = """
+function typeText(text, element) {
+    let index = 0;
+    element.innerHTML = '';
+    function type() {
+        if (index < text.length) {
+            element.innerHTML += text[index];
+            index++;
+            setTimeout(type, 50);  // Adjust typing speed here
+        } else {
+            element.innerHTML += '<span class="blinking-cursor"></span>';
+        }
+    }
+    type();
+}
+"""
+
 iface = gr.Interface(
     fn=gradio_generate,
     inputs=[
@@ -162,9 +209,12 @@ iface = gr.Interface(
         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 Text Generator",
-    description="Enter a prompt and adjust parameters to generate text using a fine-tuned GPT model."
+    outputs=gr.HTML(label="Generated Text"),
+    title="Animated GPT Text Generator",
+    description="Enter a prompt and adjust parameters to generate text using a fine-tuned GPT model.",
+    css=custom_css,
+    js=js_code,
+    live=True
 )
 
 # Launch the app