Update generate_transcript.py

generate_transcript.py

@@ -10,12 +10,12 @@ import warnings
 warnings.filterwarnings('ignore')
 
 
-class TranscriptGenerator:
+class TranscriptProcessor:
     """
-    A class to generate a conversational podcast transcript from cleaned text.
+    A class to generate and rewrite podcast-style transcripts using a specified language model.
     """
 
-    def __init__(self, text_file_path, model_name="meta-llama/Llama-3.1-70B-Instruct"):
+    def __init__(self, text_file_path, model_name="meta-llama/Llama-3.1-8B-Instruct"):
         """
         Initialize with the path to the cleaned text file and the model name.
         
@@ -24,7 +24,8 @@ class TranscriptGenerator:
             model_name (str): Name of the language model to use.
         """
         self.text_file_path = text_file_path
-        self.output_path = './resources/data.pkl'
+        self.transcript_output_path = './resources/data.pkl'
+        self.tts_output_path = './resources/podcast_ready_data.pkl'
         self.model_name = model_name
         self.accelerator = Accelerator()
         self.model = transformers.pipeline(
@@ -33,19 +34,23 @@ class TranscriptGenerator:
             model_kwargs={"torch_dtype": torch.bfloat16},
             device_map="auto"
         )
-        self.system_prompt = """
-        You are a world-class podcast writer, you have worked as a ghost writer for Joe Rogan, Lex Fridman, Ben Shapiro, Tim Ferris. 
-        We are in an alternate universe where actually you have been writing every line they say and they just stream it into their brains.
-
-        Your job is to write word by word, even "umm, hmmm, right" interruptions by the second speaker based on the PDF upload. 
-        Keep it extremely engaging, with realistic anecdotes, tangents, and interruptions.
+        self.transcript_prompt = """
+        You are a world-class podcast writer, working as a ghost writer for top podcast hosts.
+        You will write the dialogue with engaging interruptions, anecdotes, and curiosity-led questions.
+        
+        Speaker 1: Leads the conversation. Speaker 2: Asks follow-up questions and reacts with expressions.
+        
+        ALWAYS START WITH SPEAKER 1: STRICTLY THE DIALOGUES.
+        """
+        
+        self.rewrite_prompt = """
+        You are an international oscar-winning screenwriter creating a refined script for TTS.
         
-        Speaker 1: Leads and teaches. Speaker 2: Asks follow-up questions, gets excited or confused.
+        Speaker 1: Teaches with anecdotes; Speaker 2: Reacts with expressions like "umm," "hmm," [sigh].
         
-        ALWAYS START YOUR RESPONSE DIRECTLY WITH SPEAKER 1:
-        STRICTLY THE DIALOGUES.
+        Return the response as a list of tuples only, with no extra formatting.
         """
-    
+
     def load_text(self):
         """
         Reads the cleaned text file and returns its content.
@@ -77,7 +82,7 @@ class TranscriptGenerator:
             return None
         
         messages = [
-            {"role": "system", "content": self.system_prompt},
+            {"role": "system", "content": self.transcript_prompt},
             {"role": "user", "content": input_text}
         ]
         
@@ -90,7 +95,37 @@ class TranscriptGenerator:
         transcript = output[0]["generated_text"]
         
         # Save the transcript as a pickle file
-        with open(self.output_path, 'wb') as f:
+        with open(self.transcript_output_path, 'wb') as f:
             pickle.dump(transcript, f)
         
-        return self.output_path
+        return self.transcript_output_path
+
+    def rewrite_transcript(self):
+        """
+        Refines the transcript for TTS, adding expressive elements and saving as a list of tuples.
+        
+        Returns:
+            str: Path to the file where the TTS-ready transcript is saved.
+        """
+        # Load the initial generated transcript
+        with open(self.transcript_output_path, 'rb') as file:
+            input_transcript = pickle.load(file)
+        
+        messages = [
+            {"role": "system", "content": self.rewrite_prompt},
+            {"role": "user", "content": input_transcript}
+        ]
+        
+        output = self.model(
+            messages,
+            max_new_tokens=8126,
+            temperature=1
+        )
+        
+        rewritten_transcript = output[0]["generated_text"]
+        
+        # Save the rewritten transcript as a pickle file
+        with open(self.tts_output_path, 'wb') as f:
+            pickle.dump(rewritten_transcript, f)
+        
+        return self.tts_output_path