Upload app.py

app.py

@@ -220,8 +220,8 @@ print(f"Model ready on {device}")
 enc = tiktoken.get_encoding('gpt2')
 
 
-def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
-    """Generate text from prompt"""
+def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50, top_p=0.9, repetition_penalty=1.1):
+    """Generate text from prompt with improved sampling"""
     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)."
@@ -232,6 +232,8 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
         
         temperature = max(0.1, min(2.0, temperature))  # Clamp temperature
         top_k = max(1, min(100, int(top_k)))  # Clamp top_k
+        top_p = max(0.1, min(1.0, float(top_p)))  # Clamp top_p (nucleus sampling)
+        repetition_penalty = max(1.0, min(1.5, float(repetition_penalty)))  # Clamp repetition penalty
         max_new_tokens = max(1, min(200, int(max_new_tokens)))  # Clamp max tokens
         
         # Encode prompt
@@ -241,34 +243,73 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
         
         tokens = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)
         
-        # Generate
+        # Generate with improved sampling strategy
         with torch.no_grad():
+            # Track recent tokens for repetition penalty
+            recent_tokens = set()
+            
             for i in range(max_new_tokens):
                 # Forward pass
                 logits, _ = model(tokens)
-                logits = logits[:, -1, :] / max(temperature, 0.1)  # Avoid division by zero
+                logits = logits[:, -1, :] / max(temperature, 0.1)  # Apply temperature
                 
-                # 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
+                # Apply repetition penalty to reduce loops
+                if repetition_penalty > 1.0 and len(recent_tokens) > 0:
+                    for token_id in recent_tokens:
+                        if logits[0, token_id] > 0:
+                            logits[0, token_id] /= repetition_penalty
+                        else:
+                            logits[0, token_id] *= repetition_penalty
                 
-                # Sample from distribution
+                # Convert to probabilities
                 probs = F.softmax(logits, dim=-1)
                 
-                # Avoid NaN
-                if torch.isnan(probs).any():
+                # Apply top-p (nucleus) sampling first - often better than just top-k
+                if top_p < 1.0:
+                    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+                    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                    
+                    # Remove tokens with cumulative probability above threshold
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    # Keep at least one token
+                    sorted_indices_to_remove[..., 0] = False
+                    
+                    # Create mask
+                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    probs[indices_to_remove] = 0
+                    
+                    # Renormalize
+                    probs = probs / probs.sum()
+                
+                # Apply top-k filtering (after top-p for better quality)
+                if top_k < logits.size(-1):
+                    topk_probs, topk_indices = torch.topk(probs, top_k, dim=-1)
+                    # Create filtered probabilities
+                    filtered_probs = torch.zeros_like(probs)
+                    filtered_probs.scatter_(-1, topk_indices, topk_probs)
+                    # Renormalize
+                    filtered_probs = filtered_probs / filtered_probs.sum()
+                    probs = filtered_probs
+                
+                # Avoid NaN or zero probabilities
+                if torch.isnan(probs).any() or (probs.sum() == 0):
                     probs = torch.ones_like(probs) / probs.size(-1)
                 
+                # Sample from distribution
                 next_token = torch.multinomial(probs, 1)
                 
+                # Update recent tokens for repetition penalty (keep last 20 tokens)
+                token_id = next_token.item()
+                recent_tokens.add(token_id)
+                if len(recent_tokens) > 20:
+                    # Remove oldest tokens (simple approach: keep last 20)
+                    recent_tokens = set(list(recent_tokens)[-20:])
+                
                 # Append to sequence
                 tokens = torch.cat([tokens, next_token], dim=1)
                 
-                # Stop if we hit max length
+                # Early stopping: stop if we generate end-of-text token (if present)
+                # For GPT-2 tokenizer, we can check for certain patterns
                 if tokens.size(1) >= config.block_size:
                     break
         
@@ -278,22 +319,102 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
         # Post-process to fix spacing issues (common with BPE tokenizers)
         import re
         
+        # Fix 0: Remove the prompt from the beginning if it appears as a speaker name
+        # This handles cases where user enters "Romeo and Juliet" and model treats it as speaker
+        prompt_lower = prompt.lower().strip()
+        generated_lower = generated_text.lower()
+        
+        # If prompt appears at the very start and looks like it was treated as a speaker
+        if generated_lower.startswith(prompt_lower):
+            # Check if it's followed by a newline (speaker format) or dialogue
+            prompt_len = len(prompt)
+            if len(generated_text) > prompt_len:
+                next_chars = generated_text[prompt_len:prompt_len+5].strip()
+                # If prompt is followed by newline or colon-like pattern, it was treated as speaker
+                if not next_chars or ':' in next_chars or '\n' in generated_text[prompt_len:prompt_len+5]:
+                    # Remove the prompt from output (it's the input, not part of generated story)
+                    generated_text = generated_text[len(prompt):].strip()
+                    # Remove leading newlines/colons
+                    generated_text = re.sub(r'^[\s:]+', '', generated_text)
+                    
+                    # Check if the first line after removal is orphaned dialogue (no speaker)
+                    lines = generated_text.split('\n')
+                    if lines and lines[0].strip():
+                        first_line = lines[0].strip()
+                        # If first line is not a speaker name and looks like dialogue, add a speaker
+                        if not re.match(r'^([A-Z][A-Z\s]+?):\s*$', first_line):
+                            # Check if it's dialogue-like (starts with capital, has punctuation)
+                            if re.match(r'^[A-Z]', first_line) and ('.' in first_line or ',' in first_line or '!' in first_line or '?' in first_line):
+                                # Add a generic speaker name based on the prompt context
+                                # For story prompts like "Romeo and Juliet", use a character from the prompt
+                                prompt_words = [w.capitalize() for w in prompt_lower.split() if len(w) > 2]
+                                if len(prompt_words) >= 2:
+                                    # Use first significant word as speaker (e.g., "Romeo" from "Romeo and Juliet")
+                                    speaker_name = prompt_words[0].upper()
+                                else:
+                                    # Generic speaker
+                                    speaker_name = "NARRATOR"
+                                
+                                # Add speaker before the dialogue
+                                generated_text = f"{speaker_name}:\n{first_line}\n" + '\n'.join(lines[1:]) if len(lines) > 1 else f"{speaker_name}:\n{first_line}"
+        
         # Fix 1: lowercase followed by uppercase (e.g., "perpetualWith" -> "perpetual With")
         generated_text = re.sub(r'([a-z])([A-Z])', r'\1 \2', generated_text)
         
         # Fix 2: Common word boundaries that got merged (e.g., "perpetualwith" -> "perpetual with")
         # Add space before common words that might have been merged
-        common_words = ['with', 'the', 'and', 'that', 'this', 'have', 'from', 'not', 'but', 'for', 'are', 'was', 'were', 'been', 'will', 'shall', 'would', 'could', 'should']
+        common_words = ['with', 'the', 'and', 'that', 'this', 'have', 'from', 'not', 'but', 'for', 'are', 'was', 'were', 'been', 'will', 'shall', 'would', 'could', 'should', 'be', 'your', 'you', 'our', 'my', 'his', 'her', 'their', 'him', 'them']
         for word in common_words:
             # Only add space if it's not already separated and follows a lowercase letter
             pattern = r'([a-z])(' + word + r'\b)'
             generated_text = re.sub(pattern, r'\1 \2', generated_text, flags=re.IGNORECASE)
         
+        # Fix 2b: Fix contractions that got merged (e.g., "You'llbe" -> "You'll be")
+        # Add space after contractions before lowercase words
+        contractions = ["'ll", "'ve", "'re", "'d", "'t", "'s", "'m"]
+        for contraction in contractions:
+            # Pattern: contraction followed by lowercase letter (e.g., "You'llbe" -> "You'll be")
+            pattern = r"(" + re.escape(contraction) + r")([a-z])"
+            generated_text = re.sub(pattern, r'\1 \2', generated_text, flags=re.IGNORECASE)
+        
         # Fix 3: Add space before character names (all caps words)
         generated_text = re.sub(r'([a-z])([A-Z]{2,})', r'\1 \2', generated_text)
         
-        # Fix 4: Remove duplicate speaker names (e.g., "LEONTES:\n...\nLEONTES:" -> keep only first)
-        # More aggressive: remove same speaker if it appears within 5 lines
+        # Fix 3b: Normalize speaker names (e.g., "Romeo and juliet" -> "ROMEO AND JULIET:")
+        # Handle mixed case speaker names that should be all caps
+        lines = generated_text.split('\n')
+        normalized_lines = []
+        for i, line in enumerate(lines):
+            line_stripped = line.strip()
+            
+            # Check if line is a potential speaker name (title case or mixed case, 2+ words)
+            # Pattern: "Romeo and juliet", "Romeo And Juliet", etc.
+            speaker_pattern = r'^([A-Z][a-z]+(?:\s+[a-zA-Z]+)+)\s*:?\s*$'
+            match = re.match(speaker_pattern, line_stripped)
+            
+            if match:
+                # Check if next line is dialogue (not another speaker)
+                is_speaker = False
+                if i + 1 < len(lines):
+                    next_line = lines[i + 1].strip()
+                    # If next line is not empty and not a speaker name, this is likely a speaker
+                    if next_line and not re.match(r'^([A-Z][A-Z\s]+?):\s*$', next_line):
+                        is_speaker = True
+                elif i == 0:  # First line is likely a speaker if it matches pattern
+                    is_speaker = True
+                
+                if is_speaker:
+                    # Convert to all caps and ensure colon
+                    speaker_name = match.group(1).upper()
+                    normalized_lines.append(speaker_name + ':')
+                    continue
+            
+            normalized_lines.append(line)
+        
+        generated_text = '\n'.join(normalized_lines)
+        
+        # Fix 4: Remove duplicate speaker names (e.g., "EDWARD IV:\n...\nEDWARD IV:" -> keep only first)
+        # More aggressive: remove same speaker if it appears within 3 lines (tighter window)
         lines = generated_text.split('\n')
         cleaned_lines = []
         speaker_history = []  # Track recent speakers with their line numbers
@@ -306,9 +427,9 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
             if speaker_match:
                 speaker = speaker_match.group(1).strip()
                 
-                # Check if this speaker appeared recently (within last 5 lines)
+                # Check if this speaker appeared recently (within last 3 lines - more aggressive)
                 recent_speaker = False
-                for hist_speaker, hist_line_num in speaker_history[-5:]:
+                for hist_speaker, hist_line_num in speaker_history[-3:]:
                     if speaker == hist_speaker:
                         recent_speaker = True
                         break
@@ -329,11 +450,47 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.8, top_k=50):
         
         generated_text = '\n'.join(cleaned_lines)
         
-        # Fix 5: Remove multiple empty lines between speaker and dialogue
+        # Fix 5: Remove speaker names with no dialogue (e.g., "KING:\nEDWARD IV:" -> "EDWARD IV:")
+        # A speaker name should be followed by actual dialogue, not immediately by another speaker
+        lines = generated_text.split('\n')
+        final_lines = []
+        
+        for i, line in enumerate(lines):
+            line_stripped = line.strip()
+            speaker_match = re.match(r'^([A-Z][A-Z\s]+?):\s*$', line_stripped)
+            
+            if speaker_match:
+                # Check if next non-empty line is another speaker (meaning this speaker has no dialogue)
+                has_dialogue = False
+                for j in range(i + 1, min(i + 3, len(lines))):  # Check next 3 lines (more aggressive)
+                    next_line = lines[j].strip()
+                    if not next_line:  # Skip empty lines
+                        continue
+                    # If next non-empty line is NOT a speaker, we have dialogue
+                    if not re.match(r'^([A-Z][A-Z\s]+?):\s*$', next_line):
+                        has_dialogue = True
+                        break
+                    # If next non-empty line IS a speaker, this speaker has no dialogue
+                    else:
+                        # This speaker has no dialogue - skip it
+                        break
+                
+                if not has_dialogue:
+                    # This speaker has no dialogue, skip it
+                    continue
+            
+            final_lines.append(line)
+        
+        generated_text = '\n'.join(final_lines)
+        
+        # Fix 5b: Fix merged text issues (e.g., "You?A:" -> "You? A:")
+        # Add space after question/exclamation marks before capital letters
+        generated_text = re.sub(r'([?!])([A-Z])', r'\1 \2', generated_text)
+        
+        # Fix 6: Remove multiple empty lines between speaker and dialogue
         generated_text = re.sub(r'([A-Z][A-Z\s]+?):\s*\n\s*\n+', r'\1:\n', generated_text)
         
-        # Fix 6: Remove any remaining consecutive duplicate speakers (final cleanup)
-        # Pattern: Same speaker name appearing on consecutive lines (with optional whitespace)
+        # Fix 7: Remove any remaining consecutive duplicate speakers (final cleanup)
         generated_text = re.sub(
             r'^([A-Z][A-Z\s]+?):\s*\n\s*\n*\1:\s*\n',
             r'\1:\n',
@@ -389,15 +546,33 @@ with gr.Blocks(title="GPT-2 124M Shakespeare Model") as demo:
                 label="Temperature",
                 minimum=0.1,
                 maximum=2.0,
-                value=0.8,
-                step=0.1
+                value=0.7,
+                step=0.1,
+                info="Lower = more focused, Higher = more creative (0.7 recommended for better coherence)"
             )
             top_k = gr.Slider(
                 label="Top-K",
                 minimum=10,
                 maximum=100,
                 value=50,
-                step=10
+                step=10,
+                info="Number of top tokens to consider"
+            )
+            top_p = gr.Slider(
+                label="Top-P (Nucleus)",
+                minimum=0.1,
+                maximum=1.0,
+                value=0.9,
+                step=0.05,
+                info="Nucleus sampling - higher = more diverse, lower = more focused (0.9 recommended)"
+            )
+            repetition_penalty = gr.Slider(
+                label="Repetition Penalty",
+                minimum=1.0,
+                maximum=1.5,
+                value=1.1,
+                step=0.05,
+                info="Penalize repeated tokens - higher = less repetition (1.1 recommended)"
             )
             generate_btn = gr.Button("Generate", variant="primary")
         
@@ -433,7 +608,7 @@ with gr.Blocks(title="GPT-2 124M Shakespeare Model") as demo:
     
     generate_btn.click(
         fn=generate_text,
-        inputs=[prompt_input, max_tokens, temperature, top_k],
+        inputs=[prompt_input, max_tokens, temperature, top_k, top_p, repetition_penalty],
         outputs=output
     )