Upload app.py

app.py

@@ -367,7 +367,10 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50, top_p=0
             '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'
+            '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'
         ]
         for word in common_words_fix:
             word_lower = word.lower()
@@ -415,15 +418,15 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50, top_p=0
         # 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", "T h is" -> "This")
+        # 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']
+        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']
         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"
+            # 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: letter space letter space ... (all letters of the word)
+                # 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]))
@@ -437,6 +440,21 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50, top_p=0
                     # 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
@@ -485,7 +503,45 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50, top_p=0
             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) and fix missing punctuation
+        # 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:"
+                    # 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)
+                    # 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)