app.py

"""
HuggingFace Spaces App for GPT-2 124M Shakespeare Model
"""
import torch
import torch.nn as nn
from torch.nn import functional as F
import tiktoken
import gradio as gr
import math
from dataclasses import dataclass


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.c_proj.NANOGPT_SCALE_INIT = 1
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.c_proj(y)
        return y


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.gelu = nn.GELU(approximate='tanh')
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
        self.c_proj.NANOGPT_SCALE_INIT = 1

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        return x


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50257
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768


class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            wpe=nn.Embedding(config.block_size, config.n_embd),
            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f=nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.config.block_size}"
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        pos_emb = self.transformer.wpe(pos)
        tok_emb = self.transformer.wte(idx)
        x = tok_emb + pos_emb
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss


# Load model
print("Loading model...")
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', 'shwethd/gpt2-shakespeare-124m')
    
    try:
        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)
                # Restore weight sharing (broken during SafeTensors conversion)
                # lm_head.weight and transformer.wte.weight should share memory
                model.transformer.wte.weight = model.lm_head.weight
                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
                )
                # PyTorch 2.6+ requires weights_only=False for custom classes
                # This is safe since we trust our own trained model
                checkpoint = torch.load(model_path, map_location=device, weights_only=False)
                
                # 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
            )
            # PyTorch 2.6+ requires weights_only=False for custom classes
            checkpoint = torch.load(model_path, map_location=device, weights_only=False)
            
            # 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...")
        try:
            # Fallback to local file
            # PyTorch 2.6+ requires weights_only=False for custom classes
            checkpoint = torch.load('model_checkpoint_final.pt', map_location=device, weights_only=False)
            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.")
except Exception as e:
    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()
print(f"Model ready on {device}")

enc = tiktoken.get_encoding('gpt2')


def generate_text(prompt, max_new_tokens=100, temperature=0.6, top_k=40, top_p=0.8, repetition_penalty=1.15):
    """
    Generate text from prompt with improved sampling.
    
    Note on Grammar Quality:
    Despite training, grammar may not be perfect due to:
    1. Model size (124M) - relatively small for perfect grammar
    2. Training loss measures token prediction, not grammatical correctness
    3. Shakespeare text has archaic grammar patterns the model learns
    4. BPE tokenization can split words, making grammar harder to learn
    
    For better grammar, try:
    - Lower temperature (0.5-0.6) for more focused, coherent text
    - Higher repetition penalty (1.15-1.2) to reduce repetitive patterns
    - Lower top_p (0.8-0.85) for more conservative 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)."
        
        # 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
        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
        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 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)  # Apply temperature
                
                # 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
                
                # Convert to probabilities
                probs = F.softmax(logits, dim=-1)
                
                # 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)
                
                # 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
        
        # Decode
        generated_text = enc.decode(tokens[0].tolist())
        
        # 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 "First Citizen:" and model repeats it
        # Normalize prompt: remove colon, strip, convert to uppercase for comparison
        prompt_normalized = prompt.strip().replace(':', '').strip().upper()
        
        # Process all lines to find and remove prompt matches
        lines = generated_text.split('\n')
        cleaned_lines = []
        prompt_removed = False
        
        for i, line in enumerate(lines):
            line_stripped = line.strip()
            
            # Skip empty lines at the start (but only if we haven't added any content yet)
            if not line_stripped:
                if not cleaned_lines:
                    continue  # Skip leading empty lines
                else:
                    cleaned_lines.append(line)  # Keep empty lines after content starts
                    continue
            
            # Normalize line for comparison (remove colon, case-insensitive)
            line_normalized = line_stripped.replace(':', '').strip().upper()
            
            # Check if this line matches the prompt (case-insensitive, allowing for colon)
            # Check if it's a speaker name format (all caps OR title case OR mixed case)
            is_speaker_line = (re.match(r'^([A-Z][A-Z\s]+?):\s*$', line_stripped) or  # All caps: "FIRST CITIZEN:"
                              re.match(r'^([A-Z][a-z]+(?:\s+[a-zA-Z]+)+):\s*$', line_stripped) or  # Title case: "First Citizen:"
                              re.match(r'^([A-Z][A-Za-z\s]+?):\s*$', line_stripped))  # Mixed case: "First Citizen:" or "FIRST Citizen:"
            
            # If this line matches the prompt (case-insensitive), remove it
            # Be more aggressive: if it matches the prompt, remove it even if pattern doesn't match exactly
            if line_normalized == prompt_normalized and not prompt_removed:
                # Additional check: if it ends with colon, it's likely a speaker name
                if line_stripped.endswith(':'):
                    # This is the prompt appearing as a speaker - skip it
                    prompt_removed = True
                    continue
                # Also remove if it's a speaker line pattern
                elif is_speaker_line:
                    prompt_removed = True
                    continue
            
            # If we've already removed the prompt, add the line
            cleaned_lines.append(line)
        
        generated_text = '\n'.join(cleaned_lines)
        
        # If after removing prompt, first line is orphaned dialogue (no speaker), handle it
        # Keep removing orphaned dialogue at the start until we find a speaker or valid content
        while generated_text.strip():
            lines = generated_text.split('\n')
            first_line = lines[0].strip() if lines else ''
            
            if not first_line:
                # Remove empty first line
                generated_text = '\n'.join(lines[1:]) if len(lines) > 1 else ''
                continue
            
            # Check if first line is a speaker name
            is_speaker = re.match(r'^([A-Z][A-Z\s]+?):\s*$', first_line) or \
                        re.match(r'^([A-Z][a-z]+(?:\s+[a-zA-Z]+)+):\s*$', first_line)
            
            if is_speaker:
                # Found a speaker, stop removing
                break
            
            # Check if it's orphaned dialogue (starts with capital, has punctuation, but no speaker)
            if re.match(r'^[A-Z]', first_line) and ('.' in first_line or ',' in first_line or '!' in first_line or '?' in first_line):
                # Remove the orphaned first line
                generated_text = '\n'.join(lines[1:]) if len(lines) > 1 else ''
            else:
                # Not clearly orphaned dialogue, stop removing
                break
        
        # Fix 1: lowercase followed by uppercase (e.g., "perpetualWith" -> "perpetual With", "AOr" -> "A Or")
        generated_text = re.sub(r'([a-z])([A-Z])', r'\1 \2', generated_text)
        # Also fix single letter + capital word (e.g., "AOr" -> "A Or")
        generated_text = re.sub(r'\b([A-Z])([A-Z][a-z]+)', r'\1 \2', generated_text)
        
        # Fix 1b: Fix spacing issues like "furt her" -> "further", "T his" -> "This", "y our" -> "your", "th at" -> "that"
        # Remove spaces in the middle of common words - MORE AGGRESSIVE matching
        common_words_fix = [
            'further', 'this', 'that', 'there', 'where', 'here', 'their', 'your', 'our', 
            'man', 'men', 'woman', 'women', 'padua', 'content', 'gentle', 'gently',
            'house', 'neck', 'car', 'made', 'lost', 'rough', 'see', 'might', 'any', 'one',
            'well', 'newly', 'too', 'him', 'her', 'them', 'they', 'the', 'and', 'but',
            'for', 'not', 'are', 'was', 'were', 'been', 'have', 'has', 'had', 'will',
            'shall', 'would', 'could', 'should', 'be', 'is', 'it', 'he', 'she', 'we',
            'you', 'me', 'my', 'his', 'hers', 'its', 'our', 'ours', 'yours', 'theirs',
            'into', 'onto', 'upon', 'within', 'without', 'through', 'though', 'although',
            'about', 'above', 'below', 'beside', 'between', 'among', 'during', 'before',
            'after', 'while', 'until', 'since', 'because', 'together', 'honour', 'honor',
            'already', 'perfect', 'soul', 'way', 'wounds', 'tears', 'raise', 'call',
            'citizens', 'senator', 'liked', 'cold', 'incold', 'incwold', 'son', 'count',
            'happen', 'happ', 'what', 'common', 'complain', 'upon', 'she', 'honour', 'honor',
            'youth', 'ports', 'impans', 'swear', 'gods', 'please', 'standing', 'tybalt',
            'sworn', 'where', 'would', 'give', 'seize', 'before', 'repair', 'lest', 'speak',
            'woman', 'gentleman', 'deed', 'better', 'virtuous', 'done', 'broke', 'art',
            'again', 'government', 'honour', 'light', 'stands', 'fly', 'mighty', 'forth',
            'turn', 'highness', 'morning', 'hence', 'enter', 'should', 'rue', 'there',
            'confess', 'suffer', 'part', 'coronured', 'eyuls', 'unto', 'until', 'grey',
            'lady', 'evils', 'eyes', 'feat', 'worn', 'sister', 'thus', 'apparent', 'blunt',
            'not', 'most', 'worthy', 'should', 'bed', 'than', 'half', 'chaste', 'sight',
            'that', 'just', 'those', 'passes', 'stuffed', 'calm', 'then', 'little', 'great',
            'secrets', 'full', 'pray', 'duke', 'songs', 'soldier', 'worthy', 'call', 'rod',
            'respect', 'drunk', 'there', 'signior', 'gremio', 'compound', 'soft', 'unvish',
            'know', 'edward'
        ]
        for word in common_words_fix:
            word_lower = word.lower()
            # Try all possible split positions
            for i in range(1, len(word_lower)):
                first_part = word_lower[:i]
                second_part = word_lower[i:]
                
                # Pattern 1: lowercase split (e.g., "furt her" -> "further", "th at" -> "that")
                # Use word boundaries but also allow punctuation/whitespace around
                pattern1 = r'\b' + re.escape(first_part) + r'\s+' + re.escape(second_part) + r'\b'
                generated_text = re.sub(pattern1, word, generated_text, flags=re.IGNORECASE)
                
                # Pattern 2: Capital first letter (e.g., "Th at" -> "That")
                pattern2 = r'\b' + re.escape(first_part.capitalize()) + r'\s+' + re.escape(second_part) + r'\b'
                generated_text = re.sub(pattern2, word.capitalize(), generated_text)
                
                # Pattern 3: All caps (e.g., "TH AT" -> "THAT")
                pattern3 = r'\b' + re.escape(first_part.upper()) + r'\s+' + re.escape(second_part.upper()) + r'\b'
                generated_text = re.sub(pattern3, word.upper(), generated_text)
                
                # Pattern 4: Mixed case - first letter capitalized (e.g., "Th at" -> "That")
                if len(first_part) > 0:
                    pattern4 = r'\b' + re.escape(first_part[0].upper() + first_part[1:]) + r'\s+' + re.escape(second_part) + r'\b'
                    generated_text = re.sub(pattern4, word.capitalize(), generated_text, flags=re.IGNORECASE)
                
                # Pattern 5: Handle multiple splits in one word (e.g., "c o u n t" -> "count")
                # This is a special case for words that got split multiple times
                if len(word_lower) > 4:  # Only for longer words
                    # Try to find pattern like "c o u n t" or "y o u r"
                    # This is more complex, so we'll handle it separately
                    pass
        
        # 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', 'be', 'your', 'you', 'our', 'my', 'his', 'her', 'their', 'him', 'them', 'to', 'of', 'in', 'on', 'at', 'as', 'is', 'it', 'he', 'she', 'we', 'they', 'an', 'a']
        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 2c: Fix double words (e.g., "but but" -> "but")
        generated_text = re.sub(r'\b(\w+)\s+\1\b', r'\1', generated_text, flags=re.IGNORECASE)
        
        # Fix 2d: Fix spacing after commas (e.g., "What,bear" -> "What, bear")
        generated_text = re.sub(r',([a-zA-Z])', r', \1', generated_text)
        
        # Fix 1c: Fix multiple splits in one word (e.g., "c o u n t" -> "count", "y o u r" -> "your", "y our" -> "your", "T h is" -> "This")
        # Handle cases where a word got split into multiple parts
        multi_split_words = ['count', 'your', 'son', 'our', 'the', 'and', 'but', 'for', 'not', 'are', 'was', 'were', 'been', 'have', 'has', 'had', 'will', 'shall', 'would', 'could', 'should', 'be', 'is', 'it', 'he', 'she', 'we', 'they', 'you', 'me', 'my', 'his', 'her', 'them', 'him', 'this', 'that', 'there', 'where', 'here', 'their', 'what', 'common', 'complain', 'upon', 'honour', 'honor', 'youth', 'ports', 'impans', 'woman', 'gentleman', 'deed', 'better', 'virtuous', 'done', 'broke', 'art', 'again', 'government', 'light', 'stands', 'fly', 'mighty', 'forth', 'turn', 'highness', 'morning', 'hence', 'enter', 'should', 'rue', 'confess', 'suffer', 'part', 'unto', 'until', 'grey', 'lady', 'evils', 'eyes', 'feat', 'worn', 'sister', 'thus', 'apparent', 'blunt', 'most', 'worthy', 'bed', 'than', 'half', 'chaste', 'sight', 'just', 'those', 'passes', 'stuffed', 'calm', 'then', 'little', 'great', 'secrets', 'full', 'pray', 'duke', 'songs', 'soldier', 'call', 'rod', 'respect', 'drunk', 'signior', 'gremio', 'compound', 'soft', 'unvish', 'know', 'edward', 'man', 'men']
        for word in multi_split_words:
            word_lower = word.lower()
            # Create pattern for word split into individual letters with spaces
            # e.g., "c o u n t" or "y o u r" or "T h is" or "Wh at" or "y our"
            if len(word_lower) > 2:
                # Pattern 1: letter space letter space ... (all letters of the word split individually)
                letters = list(word_lower)
                pattern_parts = [re.escape(letter) + r'\s+' for letter in letters[:-1]]
                pattern_parts.append(re.escape(letters[-1]))
                pattern = r'\b' + ''.join(pattern_parts) + r'\b'
                generated_text = re.sub(pattern, word, generated_text, flags=re.IGNORECASE)
                # Also handle with some capitalization (e.g., "T h is" -> "This", "Wh at" -> "What")
                pattern_cap = r'\b' + re.escape(letters[0].upper()) + r'\s+' + ''.join([re.escape(letter) + r'\s+' for letter in letters[1:-1]]) + re.escape(letters[-1]) + r'\b'
                generated_text = re.sub(pattern_cap, word.capitalize(), generated_text)
                # Handle mixed case like "Wh at" -> "What"
                if len(letters) > 2:
                    # Pattern for "Wh at" style (first two letters capitalized, rest lowercase)
                    pattern_mixed = r'\b' + re.escape(letters[0].upper()) + re.escape(letters[1]) + r'\s+' + ''.join([re.escape(letter) + r'\s+' for letter in letters[2:-1]]) + re.escape(letters[-1]) + r'\b'
                    generated_text = re.sub(pattern_mixed, word.capitalize(), generated_text, flags=re.IGNORECASE)
                
                # Pattern 2: Handle two-part splits (e.g., "y our" -> "your", "h onour" -> "honour")
                # Try all possible two-part splits
                for split_pos in range(1, len(word_lower)):
                    first_part = word_lower[:split_pos]
                    second_part = word_lower[split_pos:]
                    # Pattern: "y our" -> "your"
                    pattern_2part = r'\b' + re.escape(first_part) + r'\s+' + re.escape(second_part) + r'\b'
                    generated_text = re.sub(pattern_2part, word, generated_text, flags=re.IGNORECASE)
                    # Capitalized version: "Y our" -> "Your"
                    pattern_2part_cap = r'\b' + re.escape(first_part.capitalize()) + r'\s+' + re.escape(second_part) + r'\b'
                    generated_text = re.sub(pattern_2part_cap, word.capitalize(), generated_text)
                    # All caps: "Y OUR" -> "YOUR"
                    pattern_2part_allcap = r'\b' + re.escape(first_part.upper()) + r'\s+' + re.escape(second_part.upper()) + r'\b'
                    generated_text = re.sub(pattern_2part_allcap, word.upper(), generated_text)
        
        # Fix 2e: Fix merged words that should be separate (e.g., "himt" -> "him to", "incwold" -> "in cold")
        # Common patterns where words got merged incorrectly
        merged_fixes = [
            # Pronoun + "t" (likely "to" got merged)
            (r'\bhimt\s+', 'him to '),  # "himt me" -> "him to me"
            (r'\bhert\s+', 'her to '),  # "hert him" -> "her to him"
            (r'\bthemt\s+', 'them to '),  # "themt us" -> "them to us"
            (r'\byout\s+', 'you to '),  # "yout me" -> "you to me"
            (r'\bhimt([,.;:!?])', r'him to\1'),  # "himt," -> "him to,"
            (r'\bhert([,.;:!?])', r'her to\1'),
            (r'\bthemt([,.;:!?])', r'them to\1'),
            (r'\byout([,.;:!?])', r'you to\1'),
            # Other merged patterns
            (r'\bincwold\b', 'in cold'),  # "incwold" -> "in cold"
            (r'\bincold\b', 'in cold'),  # "incold" -> "in cold"
            (r'\blikeled\b', 'liked'),  # "likeled" -> "liked"
            (r'\bh\s+on\s+our\b', 'honour'),  # "h on our" -> "honour"
            (r'\bh\s+on\s+or\b', 'honor'),  # "h on or" -> "honor"
            (r'\bHapp\s+up\s+on\'t\b', "Happen upon't"),  # "Happ up on't" -> "Happen upon't"
            (r'\bhapp\s+up\s+on\'t\b', "happen upon't"),
            # Fix "comm on" -> "common" (if not already fixed)
            (r'\bcomm\s+on\b', 'common'),
            (r'\bComm\s+on\b', 'Common'),
            # Fix "compl a in" -> "complain" (multiple splits)
            (r'\bcompl\s+a\s+in\b', 'complain'),
            (r'\bCompl\s+a\s+in\b', 'Complain'),
            # Fix "As s he" -> "As she"
            (r'\bAs\s+s\s+he\b', 'As she'),
            (r'\bas\s+s\s+he\b', 'as she'),
            # Fix "ag a in" -> "again" (multiple splits)
            (r'\bag\s+a\s+in\b', 'again'),
            (r'\bAg\s+a\s+in\b', 'Again'),
            # Fix "UN TO" -> "UNTO" (before Fix 3c processes it)
            (r'\bUN\s+TO\b', 'UNTO'),
            (r'\bun\s+to\b', 'unto'),
            # Fix potential word issues
            (r'\bcoronured\b', 'crowned'),  # "coronured" -> "crowned"
            (r'\beyuls\b', 'evils'),  # "eyuls" -> "evils"
            # Fix "AOr" -> "A Or" or "Or" (if it's at start of sentence)
            (r'\bAOr\b', 'A Or'),
            (r'^A Or\s+', 'Or '),  # If "A Or" is at start, might just be "Or"
            # Fix "fe at" -> "feat"
            (r'\bfe\s+at\b', 'feat'),
            (r'\bFe\s+at\b', 'Feat'),
            # Fix "MORE TH AN HALF" -> "MORE THAN HALF" (but this might be dialogue, not speaker)
            (r'\bTH\s+AN\b', 'THAN'),
            (r'\bth\s+an\b', 'than'),
            # Fix "F IT" -> "FIT" (in all caps dialogue)
            (r'\bF\s+IT\b', 'FIT'),
            (r'\bf\s+it\b', 'fit'),
            (r'\bF\s+it\b', 'Fit'),
            # Fix "C A" -> "CA" (but be careful - might be part of "C A:" speaker name)
            # Actually, "C A:" should be merged to "CA:" or might be "CLARENCE:" - handle in speaker fix
            # Fix "OUCESTER" -> "GLOUCESTER" (missing "GL" prefix)
            (r'\bOUCESTER\b', 'GLOUCESTER'),
            (r'\bOucester\b', 'Gloucester'),
            # Fix "stuff'd" -> "stuffed" (if needed, but "stuff'd" is valid Shakespeare)
            # Actually, "stuff'd" is correct Shakespeare spelling, so we'll leave it
            # Fix duplicate words: "if it be it possible" -> "if it be possible"
            (r'\bif it be it\b', 'if it be'),
            (r'\bIf it be it\b', 'If it be'),
            # Fix duplicate "belike" -> remove one
            (r'\bbelike\s+that\s+you\s+were\s+right\s+gentle\s+exercise,\s+belike\b', 'belike that you were right gentle exercise'),
            (r'\bBelike\s+that\s+you\s+were\s+right\s+gentle\s+exercise,\s+belike\b', 'Belike that you were right gentle exercise'),
        ]
        for pattern, replacement in merged_fixes:
            generated_text = re.sub(pattern, replacement, generated_text, flags=re.IGNORECASE)
        
        # Fix 2f: Fix "content on" - this is likely two separate words, but ensure proper spacing
        generated_text = re.sub(r'\bcontenton\b', 'content on', generated_text, flags=re.IGNORECASE)
        
        # Fix 2g: Fix "toget her" -> "together"
        generated_text = re.sub(r'\btoget\s+her\b', 'together', 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: Fix split speaker names (e.g., "ALL ANC A:" -> "ALLANCA:", "GENTLEM AN:" -> "GENTLEMAN:")
        # Pattern: All caps words separated by spaces ending with colon (likely split speaker name)
        # First, try to merge split speaker names: "ALL ANC A:" -> "ALLANCA:", "GENTLEM AN:" -> "GENTLEMAN:"
        # But be careful - some speaker names might legitimately have spaces (e.g., "FIRST CITIZEN:")
        lines = generated_text.split('\n')
        fixed_lines = []
        for line in lines:
            line_stripped = line.strip()
            # Check if line looks like a split speaker name (all caps, has spaces, ends with colon)
            # Pattern 1: Multiple all-caps words with spaces: "ALL ANC A:" or "GENTLEM AN:"
            if re.match(r'^([A-Z]+\s+[A-Z]+\s*[A-Z]*):\s*$', line_stripped):
                # Check if it's a known multi-word speaker name (keep those)
                known_multi_word_speakers = ['FIRST CITIZEN', 'SECOND CITIZEN', 'THIRD CITIZEN', 
                                            'FIRST GENTLEMAN', 'SECOND GENTLEMAN', 'THIRD GENTLEMAN',
                                            'FIRST SERVANT', 'SECOND SERVANT', 'LADY MACBETH',
                                            'KING HENRY', 'PRINCE HAMLET', 'DUKE VINCENTIO']
                is_known = False
                for known in known_multi_word_speakers:
                    if known in line_stripped.upper():
                        is_known = True
                        break
                
                if not is_known:
                    # Try to merge: "ALL ANC A:" -> "ALLANCA:", "GENTLEM AN:" -> "GENTLEMAN:", "C A:" -> "CA:" or "CLARENCE:"
                    # Remove spaces between all-caps words before colon
                    merged = re.sub(r'([A-Z]+)\s+([A-Z]+)\s*([A-Z]*):', r'\1\2\3:', line_stripped)
                    
                    # Special case: "C A:" might be "CLARENCE:" - check if it's a known pattern
                    if re.match(r'^C\s+A:\s*$', line_stripped):
                        # Check context - if it's near "Clarence" or "Sir Clarence", it's likely "CLARENCE:"
                        # For now, merge to "CA:" and let it be handled as a potential speaker
                        merged = 'CLARENCE:'
                    
                    # Only use merged if it makes sense (not too long, looks like a word)
                    if len(merged) < 30:  # Reasonable speaker name length
                        fixed_lines.append(merged)
                    else:
                        fixed_lines.append(line)
                else:
                    # Keep known multi-word speaker names as is
                    fixed_lines.append(line)
            else:
                fixed_lines.append(line)
        generated_text = '\n'.join(fixed_lines)
        
        # Fix 3b: Add space before character names (all caps words) and fix missing punctuation
        # First, fix cases like "Barn MENENIUS:" -> "Barn. MENENIUS:" or "Barn, MENENIUS:"
        # Pattern: lowercase word followed immediately by all-caps speaker name
        generated_text = re.sub(r'([a-z]+)([A-Z]{2,}):', r'\1. \2:', generated_text)
        # Then add space before character names
        generated_text = re.sub(r'([a-z])([A-Z]{2,})', r'\1 \2', generated_text)
        
        # Fix 3b: Normalize speaker names (e.g., "Romeo and juliet" -> "ROMEO AND JULIET:")
        # Handle mixed case speaker names that should be all caps
        # Also handle "First Citizen:" -> "FIRST CITIZEN:"
        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", "First Citizen", etc.
            speaker_pattern = r'^([A-Z][a-z]+(?:\s+[a-zA-Z]+)+)\s*:?\s*$'
            match = re.match(speaker_pattern, line_stripped)
            
            # Also check for all-caps speaker names (already normalized)
            all_caps_speaker = re.match(r'^([A-Z][A-Z\s]+?):\s*$', 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
            elif all_caps_speaker:
                # Already all caps, just ensure it has colon
                speaker_name = all_caps_speaker.group(1).strip()
                if not line_stripped.endswith(':'):
                    normalized_lines.append(speaker_name + ':')
                else:
                    normalized_lines.append(line)
                continue
            
            normalized_lines.append(line)
        
        generated_text = '\n'.join(normalized_lines)
        
        # Fix 0b: Remove prompt again after normalization (in case it was normalized to all caps)
        # This handles cases where "First Citizen:" was normalized to "FIRST CITIZEN:"
        prompt_normalized = prompt.strip().replace(':', '').strip().upper()
        lines = generated_text.split('\n')
        cleaned_lines_after_norm = []
        prompt_removed_after_norm = False
        
        for i, line in enumerate(lines):
            line_stripped = line.strip()
            
            # Skip empty lines at the start
            if not line_stripped and not cleaned_lines_after_norm:
                continue
            
            # Normalize line for comparison (remove colon, case-insensitive)
            line_normalized = line_stripped.replace(':', '').strip().upper()
            
            # Check if this line matches the prompt (case-insensitive, allowing for colon)
            # Also check if it's a speaker name format (all caps after normalization)
            is_speaker_line = re.match(r'^([A-Z][A-Z\s]+?):\s*$', line_stripped)
            
            if is_speaker_line and line_normalized == prompt_normalized and not prompt_removed_after_norm:
                # This is the prompt appearing as a speaker - skip it
                prompt_removed_after_norm = True
                continue
            
            # If we've already removed the prompt, add the line
            cleaned_lines_after_norm.append(line)
        
        generated_text = '\n'.join(cleaned_lines_after_norm)
        
        # Fix 3c: Fix dialogue that was incorrectly formatted as speaker names
        # Pattern: All caps lines ending with colon that are actually dialogue (not speakers)
        # Examples: "HENCE ARE YOUR HONOUR TO ENTER:" -> "HENCE ARE YOUR HONOUR TO ENTER."
        #           "THERE SHOULD RUE:" -> "THERE SHOULD RUE."
        #           "UN TO THE LADY GREY:" -> "UNTO THE LADY GREY."
        # These are usually long phrases (3+ words) that don't look like character names
        lines = generated_text.split('\n')
        fixed_dialogue_lines = []
        # Known speaker names (keep these as speakers)
        known_speakers = ['BAPTISTA', 'GLOUCESTER', 'CLARENCE', 'ROMEO', 'JULIET', 'HAMLET', 'MACBETH', 
                         'KING', 'QUEEN', 'DUKE', 'PRINCE', 'LADY', 'FIRST', 'SECOND', 'THIRD',
                         'CITIZEN', 'GENTLEMAN', 'SERVANT', 'MENENIUS', 'COMINIUS', 'CORIOLANUS',
                         'VINCENTIO', 'ANGELO', 'ISABELLA', 'OTHELLO', 'DESDEMONA', 'IAGO']
        
        for i, line in enumerate(lines):
            line_stripped = line.strip()
            # Check if line looks like all-caps speaker but is actually dialogue
            # Pattern: All caps, ends with colon
            if re.match(r'^([A-Z][A-Z\s]+?):\s*$', line_stripped):
                words = line_stripped.split()
                speaker_name = words[0] if words else ''
                
                # Check if it's a known speaker name (1-2 words, known name)
                is_known_speaker = (len(words) <= 2 and speaker_name in known_speakers) or \
                                  (len(words) == 2 and words[0] in ['FIRST', 'SECOND', 'THIRD'] and words[1] in ['CITIZEN', 'GENTLEMAN', 'SERVANT'])
                
                if is_known_speaker:
                    # Keep as speaker name
                    fixed_dialogue_lines.append(line)
                # If it has 3+ words, it's likely dialogue, not a speaker name
                elif len(words) >= 3:
                    # Convert colon to period (dialogue ending)
                    dialogue = line_stripped[:-1] + '.'  # Remove colon, add period
                    fixed_dialogue_lines.append(dialogue)
                # Also check if it contains common dialogue words (not speaker names)
                elif any(word in ['ARE', 'YOUR', 'HONOUR', 'TO', 'ENTER', 'SHOULD', 'RUE', 'THE', 'GREY', 'HENCE', 'THERE', 'UN', 'UNTIL', 'UNTO', 'MORE', 'THAN', 'HALF', 'TH', 'AN'] for word in words):
                    # Likely dialogue, not speaker
                    dialogue = line_stripped[:-1] + '.'  # Remove colon, add period
                    fixed_dialogue_lines.append(dialogue)
                # Special case: Single letter "A:" is likely dialogue or incomplete, not a speaker
                elif len(words) == 1 and words[0] == 'A':
                    # Convert to dialogue
                    fixed_dialogue_lines.append('A.')
                # Special case: "MORE THAN HALF:" is dialogue, not speaker
                elif 'MORE' in words and 'THAN' in words:
                    dialogue = line_stripped[:-1] + '.'  # Remove colon, add period
                    fixed_dialogue_lines.append(dialogue)
                else:
                    # Keep as speaker name (might be a short unknown character name)
                    fixed_dialogue_lines.append(line)
            else:
                fixed_dialogue_lines.append(line)
        
        generated_text = '\n'.join(fixed_dialogue_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 5 lines (expanded window for empty lines)
        # Also handle case-insensitive duplicates (e.g., "First Citizen:" and "FIRST CITIZEN:")
        lines = generated_text.split('\n')
        cleaned_lines = []
        speaker_history = []  # Track recent speakers with their line numbers (case-insensitive)
        
        for i, line in enumerate(lines):
            line_stripped = line.strip()
            # Check if this line is a speaker name (all caps after normalization)
            speaker_match = re.match(r'^([A-Z][A-Z\s]+?):\s*$', line_stripped)
            
            if speaker_match:
                speaker = speaker_match.group(1).strip()
                speaker_upper = speaker.upper()  # For case-insensitive comparison
                
                # Check if this speaker appeared recently (within last 5 lines - expanded for empty lines)
                # Check both exact match and case-insensitive match
                recent_speaker = False
                for hist_speaker, hist_line_num in speaker_history[-5:]:  # Check last 5 speakers
                    hist_speaker_upper = hist_speaker.upper()
                    if speaker == hist_speaker or speaker_upper == hist_speaker_upper:
                        recent_speaker = True
                        break
                
                if recent_speaker:
                    # Skip this duplicate speaker
                    continue
                
                # Add to history (store uppercase version for consistent comparison)
                speaker_history.append((speaker_upper, i))
                # Keep only last 15 speakers in history (expanded)
                if len(speaker_history) > 15:
                    speaker_history.pop(0)
                
                cleaned_lines.append(line)
            else:
                cleaned_lines.append(line)
        
        generated_text = '\n'.join(cleaned_lines)
        
        # Fix 5: Remove speaker names with no dialogue (e.g., "KING:\nEDWARD IV:" -> "EDWARD IV:", "First Citizen:\n\nCLARENCE:" -> "CLARENCE:")
        # A speaker name should be followed by actual dialogue, not immediately by another speaker or empty lines
        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 or if there's no dialogue at all
                has_dialogue = False
                # Check up to 5 lines ahead (more generous to catch dialogue)
                for j in range(i + 1, min(i + 6, len(lines))):
                    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
                    elif re.match(r'^([A-Z][A-Z\s]+?):\s*$', next_line):
                        # This speaker has no dialogue - skip it
                        has_dialogue = False
                        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 7: Remove any remaining consecutive duplicate speakers (final cleanup)
        # Handle both exact duplicates and case-insensitive duplicates
        # This handles cases like "FIRST CITIZEN:\n\nFIRST CITIZEN:" -> "FIRST CITIZEN:"
        lines = generated_text.split('\n')
        final_cleaned_lines = []
        last_speaker_upper = None
        
        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:
                speaker = speaker_match.group(1).strip()
                speaker_upper = speaker.upper()
                
                # If this is the same speaker as the last one (case-insensitive), skip it
                if speaker_upper == last_speaker_upper:
                    continue
                
                last_speaker_upper = speaker_upper
                final_cleaned_lines.append(line)
            else:
                # Reset speaker tracking on non-speaker lines (but keep last_speaker for nearby duplicates)
                # Only reset if we have substantial content (not just empty lines)
                if line_stripped:  # Non-empty line
                    # Keep last_speaker for a few lines in case of empty lines
                    pass
                final_cleaned_lines.append(line)
        
        generated_text = '\n'.join(final_cleaned_lines)
        
        # Fix 7b: Remove duplicate words in sentences (e.g., "if it be it possible" -> "if it be possible")
        # Pattern: word followed by same word (case-insensitive)
        # But be careful not to remove valid repetitions like "very very" or "more more"
        # Only remove common function words that shouldn't repeat
        duplicate_word_patterns = [
            (r'\b(it)\s+\1\b', r'\1'),  # "it it" -> "it"
            (r'\b(the)\s+\1\b', r'\1'),  # "the the" -> "the"
            (r'\b(a)\s+\1\b', r'\1'),  # "a a" -> "a"
            (r'\b(an)\s+\1\b', r'\1'),  # "an an" -> "an"
            (r'\b(is)\s+\1\b', r'\1'),  # "is is" -> "is"
            (r'\b(was)\s+\1\b', r'\1'),  # "was was" -> "was"
            (r'\b(are)\s+\1\b', r'\1'),  # "are are" -> "are"
            (r'\b(be)\s+\1\b', r'\1'),  # "be be" -> "be"
        ]
        for pattern, replacement in duplicate_word_patterns:
            generated_text = re.sub(pattern, replacement, generated_text, flags=re.IGNORECASE)
        
        # Fix 8: Handle incomplete termination - remove incomplete words/sentences at the end
        # This happens when the model hits the token limit mid-generation
        if generated_text.strip():
            # Remove incomplete word at the end (word that doesn't end with punctuation or space)
            # Pattern: ends with a word that has no trailing punctuation/space
            # But keep if it ends with proper punctuation (. ! ? , ; :)
            lines = generated_text.split('\n')
            if lines:
                last_line = lines[-1].strip()
                
                # If last line doesn't end with punctuation and is not a speaker name
                if last_line and not re.match(r'^([A-Z][A-Z\s]+?):\s*$', last_line):
                    # Check if it ends with incomplete word (no punctuation, not a complete sentence)
                    # Remove if it ends with a word that looks incomplete
                    # Pattern: ends with word that has no punctuation
                    if not re.search(r'[.!?,;:]$', last_line):
                        # Check if the last "word" is very short (likely incomplete)
                        # Or if it's a single character/letter (likely cut off)
                        words = last_line.split()
                        if words:
                            last_word = words[-1]
                            # If last word is very short (1-2 chars) and not punctuation, likely incomplete
                            if len(last_word) <= 2 and last_word.isalpha():
                                # Remove the incomplete last word
                                lines[-1] = ' '.join(words[:-1]) if len(words) > 1 else ''
                            # If last word doesn't end with punctuation and line is short, might be incomplete
                            elif len(last_line) < 20 and not last_word.endswith(('.', '!', '?', ',', ';', ':')):
                                # Check if removing last word makes sense
                                # Only remove if it's clearly incomplete (very short word)
                                if len(last_word) < 4:
                                    lines[-1] = ' '.join(words[:-1]) if len(words) > 1 else ''
                    
                    # If after processing, last line is empty or just whitespace, remove it
                    if not lines[-1].strip():
                        lines = lines[:-1]
                
                # Reconstruct text
                generated_text = '\n'.join(lines)
                
                # Final check: if text doesn't end with punctuation and is not a speaker, 
                # try to find the last complete sentence
                # BUT: Be less aggressive - only remove if we have multiple sentences and last one is clearly incomplete
                if generated_text.strip():
                    # Find the last complete sentence (ends with . ! ?)
                    # Split by sentences
                    sentences = re.split(r'([.!?]+)', generated_text)
                    if len(sentences) > 3:  # Only if we have at least 2 complete sentences
                        # Reconstruct, keeping only complete sentences
                        complete_text = ''
                        for i in range(0, len(sentences) - 1, 2):
                            if i + 1 < len(sentences):
                                complete_text += sentences[i] + sentences[i + 1]
                        # If we have complete sentences, use them; otherwise keep original
                        if complete_text.strip():
                            # But check if we removed too much (more than 30% of text must remain)
                            # AND the last sentence must be very short (likely incomplete)
                            original_len = len(generated_text.strip())
                            complete_len = len(complete_text.strip())
                            if complete_len > original_len * 0.3:
                                # Check if last sentence in original is very short (likely incomplete)
                                last_sentence = sentences[-2] if len(sentences) >= 2 else ''
                                if len(last_sentence.strip()) < 15:  # Very short last sentence
                                    generated_text = complete_text.strip()
        
        return generated_text
    except Exception as 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:
    # 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:**
    - Final Loss: 0.095127 (Target: < 0.099999) ✅
    - Model Parameters: 124.44M
    - Training Steps: 1,637
    
    **📝 Note on Grammar Quality:**
    The model learns token patterns from Shakespeare text, but grammar may not be perfect due to:
    - Model size (124M is relatively small for perfect grammar)
    - Training loss measures token prediction, not grammatical correctness
    - Shakespeare text has archaic grammar patterns
    - For better grammar, use lower temperature (0.5-0.6) and higher repetition penalty (1.15-1.2)
    
    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():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Enter your prompt here (e.g., 'First Citizen:', 'ROMEO:', 'To be or not')",
                value="First Citizen:",
                lines=3
            )
            max_tokens = gr.Slider(
                label="Max Tokens",
                minimum=50,
                maximum=200,
                value=100,
                step=10
            )
            temperature = gr.Slider(
                label="Temperature",
                minimum=0.1,
                maximum=2.0,
                value=0.6,
                step=0.1,
                info="Lower = more focused/grammatical (0.5-0.6 recommended). Higher = more creative but less coherent"
            )
            top_k = gr.Slider(
                label="Top-K",
                minimum=10,
                maximum=100,
                value=40,
                step=10,
                info="Number of top tokens to consider. Lower (30-40) = more focused, Higher = more diverse"
            )
            top_p = gr.Slider(
                label="Top-P (Nucleus)",
                minimum=0.1,
                maximum=1.0,
                value=0.8,
                step=0.05,
                info="Nucleus sampling. Lower (0.75-0.8) = more grammatical/coherent. Higher (0.9+) = more random"
            )
            repetition_penalty = gr.Slider(
                label="Repetition Penalty",
                minimum=1.0,
                maximum=1.5,
                value=1.15,
                step=0.05,
                info="Penalize repeated tokens. Higher (1.15-1.2) = better grammar/less repetition. Lower = more repetitive"
            )
            generate_btn = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(
                label="Generated Text",
                lines=10,
                interactive=True,  # Make it interactive so users can select and copy
                show_copy_button=True  # Add copy button
            )
    
    # Example prompts with suggested parameters
    gr.Markdown("### Example Prompts (Click to try - includes optimal settings)")
    examples = gr.Examples(
        examples=[
            # Format: [prompt, max_tokens, temperature, top_k, top_p, repetition_penalty]
            # Optimized for better grammar: lower temp (0.5-0.6), lower top_p (0.75-0.8), higher rep_penalty (1.15-1.2)
            ["First Citizen:", 100, 0.6, 40, 0.8, 1.15],
            ["ROMEO:", 100, 0.55, 35, 0.78, 1.2],  # Romantic - lower temp for coherence
            ["To be or not", 80, 0.5, 30, 0.75, 1.2],  # Quote - very focused for grammar
            ["HAMLET:", 100, 0.6, 40, 0.8, 1.15],
            ["MACBETH:", 100, 0.6, 40, 0.8, 1.15],
            ["JULIET:", 100, 0.55, 35, 0.78, 1.2],  # Romantic - lower temp
            ["KING:", 100, 0.6, 40, 0.8, 1.15],
            ["LADY MACBETH:", 100, 0.6, 40, 0.8, 1.15],
            ["OTHELLO:", 100, 0.6, 40, 0.8, 1.15],
            ["What light through yonder", 100, 0.55, 35, 0.78, 1.2],  # Romantic quote
            ["All the world's a stage", 100, 0.6, 40, 0.8, 1.15],  # Metaphorical
            ["Double, double toil and trouble", 80, 0.6, 40, 0.8, 1.15],  # Witches chant
            ["Friends, Romans, countrymen", 100, 0.6, 40, 0.8, 1.15],  # Speech
            ["A rose by any other name", 100, 0.55, 35, 0.78, 1.2],  # Romantic quote
        ],
        inputs=[prompt_input, max_tokens, temperature, top_k, top_p, repetition_penalty]
    )
    
    generate_btn.click(
        fn=generate_text,
        inputs=[prompt_input, max_tokens, temperature, top_k, top_p, repetition_penalty],
        outputs=output
    )
    
    gr.Markdown("""
    ---
    **Note:** The model was trained on Shakespeare text and generates text in that style.
    Generated text may not always be coherent but should follow Shakespearean patterns.
    """)

if __name__ == "__main__":
    # Don't use share=True on HuggingFace Spaces
    demo.launch()