Update app.py

app.py

@@ -149,12 +149,54 @@ class PodcastGenerator:
                     
                 add_log(f"✅ PDF extraction complete. Text length: {len(text)} characters")
                 return text.strip()
-    
+
         except Exception as e:
             error_msg = f"❌ PDF extraction failed: {str(e)}"
             add_log(error_msg)
             raise Exception(error_msg)
 
+    async def postprocess_conversation(self, raw_text: str) -> str:
+        """Run LLM again to enforce strict Speaker 1/2 format"""
+        prompt = f"""
+You are a podcast formatter.
+
+Take the following input conversation, and reformat it so that:
+- Every line begins with exactly `Speaker 1:` or `Speaker 2:` (with colon)
+- No timestamps, names, parentheses, or extra formatting
+- No blank lines
+- Do not invent or change the content
+
+Example output:
+Speaker 1: Hello and welcome.
+Speaker 2: Thanks! Glad to be here.
+
+Now format the following:
+{raw_text}
+"""
+
+        inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048
+        )
+        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=1024,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id
+            )
+
+        formatted = self.tokenizer.decode(
+            outputs[0][inputs['input_ids'].shape[1]:],
+            skip_special_tokens=True
+        )
+        return formatted.strip()
+
+
     def clean_and_validate_json(self, text: str) -> Dict:
         """Improved JSON extraction and validation - CRITICAL FIX #4"""
         add_log("🔍 Attempting to extract JSON from generated text")
@@ -379,6 +421,12 @@ Speaker 2: ...
             
             add_log(f"📝 Generated text length: {len(generated_text)} characters")
             add_log(f"🔍 Generated text preview: {generated_text[:2000]}...")
+            
+            formatted_text = await self.postprocess_conversation(generated_text)
+            add_log(f"🧼 Post-processed text:\n{formatted_text[:2000]}")
+
+            # Proceed with parsing to JSON
+            generated_text = self.conversation_to_json(formatted_text)
 
             if progress:
                 progress(0.4, "🔍 Processing generated script...")