Upload app.py

app.py

@@ -368,15 +368,31 @@ def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50, top_p=0
         generated_text = '\n'.join(cleaned_lines)
         
         # If after removing prompt, first line is orphaned dialogue (no speaker), handle it
-        if generated_text.strip():
+        # 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 ''
-            # Check if first line is orphaned dialogue (starts with capital, has punctuation, but no speaker)
-            if first_line and 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):
-                    # Just remove the orphaned first line, don't add a speaker
-                    generated_text = '\n'.join(lines[1:]) if len(lines) > 1 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)