Update app.py

app.py

@@ -1,64 +1,346 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+from dataclasses import dataclass
+from pathlib import Path
+import json
+import hashlib
+import os
+from typing import List, Tuple, Iterator
+import assemblyai as aai
+from google import generativeai
+from pydub import AudioSegment
+import asyncio
+import io
+from itertools import groupby
 
+prompt = '''
+You are an expert transcript editor. Your task is to enhance this transcript for maximum readability while maintaining the core message.
 
+IMPORTANT: Respond ONLY with the enhanced transcript. Do not include any explanations, headers, or phrases like "Here is the transcript."
+
+Note: Below you'll find an auto-generated transcript that may help with speaker identification, but focus on creating your own high-quality transcript from the audio.
+
+Think about your job as if you were transcribing an interview for a print book where the priority is the reading audience. It should just be a total pleasure to read this as a written artifact where all the flubs and repetitions and conversational artifacts and filler words and false starts are removed, where a bunch of helpful punctuation is added. It should basically read like somebody wrote it specifically for reading rather than just something somebody said extemporaneously. 
+
+Please:
+1. Fix speaker attribution errors, especially at segment boundaries. Watch for incomplete thoughts that were likely from the previous speaker.
+
+2. Optimize AGGRESSIVELY for readability over verbatim accuracy:
+   - Readability is the most important thing!!
+   - Remove ALL conversational artifacts (yeah, so, I mean, etc.)
+   - Remove ALL filler words (um, uh, like, you know)
+   - Remove false starts and self-corrections completely
+   - Remove redundant phrases and hesitations
+   - Convert any indirect or rambling responses into direct statements
+   - Break up run-on sentences into clear, concise statements
+   - Maintain natural conversation flow while prioritizing clarity and directness
+
+3. Format the output consistently:
+   - Keep the "Speaker X 00:00:00" format (no brackets, no other formatting)
+   - DO NOT change the timestamps. You're only seeing a chunk of the full transcript, which is why your 0:00:00 is not the true beginning. Keep the timestamps as they are.
+   - Add TWO line breaks between speaker/timestamp and the text
+   - Use proper punctuation and capitalization
+   - Add paragraph breaks for topic changes
+   - When you add paragraph breaks between the same speaker's remarks, no need to restate the speaker attribution
+   - Don't go more than four sentences without adding a paragraph break. Be liberal with your paragraph breaks. 
+   - Preserve distinct speaker turns
+
+Example input:
+Speaker A 00:01:15
+
+Um, yeah, so like, I've been working on this new project at work, you know? And uh, what's really interesting is that, uh, we're seeing these amazing results with the new approach we're taking. Like, it's just, you know, it's really transforming how we do things.
+
+And then, I mean, the thing is, uh, when we showed it to the client last week, they were just, you know, completely blown away by what we achieved. Like, they couldn't even believe it was the same system they had before.
+
+Example output:
+Speaker A 00:01:15
+
+I've been working on this new project at work, and we're seeing amazing results with our new approach. It's really transforming how we do things.
+
+When we showed it to the client last week, they were completely blown away by what we achieved. They couldn't believe it was the same system they had before.
+
+Enhance the following transcript, starting directly with the speaker format: 
+'''
+
+
+@dataclass
+class Utterance:
+    """A single utterance from a speaker"""
+    speaker: str
+    text: str
+    start: int
+    end: int
+
+    @property
+    def timestamp(self) -> str:
+        seconds = int(self.start // 1000)
+        hours = seconds // 3600
+        minutes = (seconds % 3600) // 60
+        seconds = seconds % 60
+        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
+
+
+class Transcriber:
+    def __init__(self, api_key: str):
+        aai.settings.api_key = api_key
+        self.cache_dir = Path("transcripts/.cache")
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def get_transcript(self, audio_path: Path) -> List[Utterance]:
+        cache_file = self.cache_dir / f"{audio_path.stem}.json"
+        
+        if cache_file.exists():
+            with open(cache_file) as f:
+                data = json.load(f)
+                if data["hash"] == self._get_file_hash(audio_path):
+                    return [
+                        Utterance(
+                            speaker=u["speaker"],
+                            text=u["text"],
+                            start=u["start"],
+                            end=u["end"]
+                        )
+                        for u in data["utterances"]
+                    ]
+
+        config = aai.TranscriptionConfig(speaker_labels=True, language_code="en")
+        transcript = aai.Transcriber().transcribe(str(audio_path), config=config)
+        
+        utterances = [
+            Utterance(
+                speaker=u.speaker,
+                text=u.text,
+                start=u.start,
+                end=u.end
+            )
+            for u in transcript.utterances
+        ]
+        
+        cache_data = {
+            "hash": self._get_file_hash(audio_path),
+            "utterances": [
+                {
+                    "speaker": u.speaker,
+                    "text": u.text,
+                    "start": u.start,
+                    "end": u.end
+                }
+                for u in utterances
+            ]
+        }
+        with open(cache_file, "w") as f:
+            json.dump(cache_data, f, indent=2)
+            
+        return utterances
+
+    def _get_file_hash(self, file_path: Path) -> str:
+        hash_md5 = hashlib.md5()
+        with open(file_path, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hash_md5.update(chunk)
+        return hash_md5.hexdigest()
+
+
+class Enhancer:
+    def __init__(self, api_key: str):
+        generativeai.configure(api_key=api_key)
+        self.model = generativeai.GenerativeModel("gemini-2.0-flash-lite-preview-02-05")
+        self.prompt = prompt
+
+    async def enhance_chunks(self, chunks: List[Tuple[str, io.BytesIO]]) -> List[str]:
+        semaphore = asyncio.Semaphore(3)
+        
+        async def process_chunk(i: int, chunk: Tuple[str, io.BytesIO]) -> str:
+            text, audio = chunk
+            async with semaphore:
+                audio.seek(0)
+                response = await self.model.generate_content_async(
+                    [self.prompt, text, {"mime_type": "audio/mp3", "data": audio.read()}]
+                )
+                return response.text
+
+        tasks = [process_chunk(i, chunk) for i, chunk in enumerate(chunks)]
+        results = await asyncio.gather(*tasks)
+        return results
+
+
+@dataclass
+class SpeakerDialogue:
+    speaker: str
+    utterances: List[Utterance]
+    
+    @property
+    def start(self) -> int:
+        return self.utterances[0].start
+    
+    @property
+    def end(self) -> int:
+        return self.utterances[-1].end
+    
+    @property
+    def timestamp(self) -> str:
+        return self.utterances[0].timestamp
+    
+    def format(self, markdown: bool = False) -> str:
+        texts = [u.text + "\n\n" for u in self.utterances]
+        combined_text = ''.join(texts).rstrip()
+        if markdown:
+            return f"**Speaker {self.speaker}** *{self.timestamp}*\n\n{combined_text}"
+        return f"Speaker {self.speaker} {self.timestamp}\n\n{combined_text}"
+
+
+def group_utterances_by_speaker(utterances: List[Utterance]) -> Iterator[SpeakerDialogue]:
+    for speaker, group in groupby(utterances, key=lambda u: u.speaker):
+        yield SpeakerDialogue(speaker=speaker, utterances=list(group))
+
+
+def estimate_tokens(text: str, chars_per_token: int = 4) -> int:
+    return (len(text) + chars_per_token - 1) // chars_per_token
+
+
+def chunk_dialogues(dialogues: Iterator[SpeakerDialogue], max_tokens: int = 2000, chars_per_token: int = 4) -> List[List[SpeakerDialogue]]:
+    chunks = []
+    current_chunk = []
+    current_text = ""
+    
+    for dialogue in dialogues:
+        formatted = dialogue.format()
+        new_text = current_text + "\n\n" + formatted if current_text else formatted
+        
+        if current_chunk and estimate_tokens(new_text, chars_per_token) > max_tokens:
+            chunks.append(current_chunk)
+            current_chunk = [dialogue]
+            current_text = formatted
+        else:
+            current_chunk.append(dialogue)
+            current_text = new_text
+    
+    if current_chunk:
+        chunks.append(current_chunk)
+    
+    return chunks
+
+
+def format_chunk(dialogues: List[SpeakerDialogue], markdown: bool = False) -> str:
+    return "\n\n".join(dialogue.format(markdown=markdown) for dialogue in dialogues)
+
+
+def prepare_audio_chunks(audio_path: Path, utterances: List[Utterance]) -> List[Tuple[str, io.BytesIO]]:
+    dialogues = group_utterances_by_speaker(utterances)
+    chunks = chunk_dialogues(dialogues)
+    
+    audio = AudioSegment.from_file(audio_path)
+    
+    prepared = []
+    for chunk in chunks:
+        segment = audio[chunk[0].start:chunk[-1].end]
+        buffer = io.BytesIO()
+        segment.export(buffer, format="mp3", parameters=["-q:a", "9"])
+        prepared.append((format_chunk(chunk, markdown=False), buffer))
+    
+    return prepared
+
+
+def apply_markdown_formatting(text: str) -> str:
+    import re
+    pattern = r"(Speaker \w+) (\d{2}:\d{2}:\d{2})"
+    return re.sub(pattern, r"**\1** *\2*", text)
+
+
+def rename_speakers(text: str, speaker_map: dict) -> str:
+    """Replace speaker labels using the provided mapping"""
+    result = text
+    for old_name, new_name in speaker_map.items():
+        # Replace both markdown and plain text formats
+        result = result.replace(f"**Speaker {old_name}**", f"**{new_name}**")
+        result = result.replace(f"Speaker {old_name}", new_name)
+    return result
+
+
+def process_audio(audio_file):
+    try:
+        # Save uploaded file with a temporary name
+        temp_path = Path("temp_audio")
+        temp_path.mkdir(exist_ok=True)
+        temp_file = temp_path / "temp_audio.mp3"
+        
+        try:
+            with open(temp_file, "wb") as f:
+                f.write(audio_file)
+
+            # Get transcript
+            transcriber = Transcriber(os.getenv("ASSEMBLYAI_API_KEY"))
+            utterances = transcriber.get_transcript(temp_file)
+            
+            # Generate original transcript
+            dialogues = list(group_utterances_by_speaker(utterances))
+            original = format_chunk(dialogues, markdown=True)
+            
+            # Show original transcript immediately
+            yield original, ""
+            
+            try:
+                # Enhance transcript
+                enhancer = Enhancer(os.getenv("GOOGLE_API_KEY"))
+                chunks = prepare_audio_chunks(temp_file, utterances)
+                enhanced = asyncio.run(enhancer.enhance_chunks(chunks))
+                
+                # Format final transcript
+                merged = "\n\n".join(chunk.strip() for chunk in enhanced)
+                merged = apply_markdown_formatting(merged)
+                
+                yield original, merged
+                
+            except Exception as e:
+                yield original, f"Error: {str(e)}"
+            
+        finally:
+            # Cleanup temp file
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+        
+    except Exception as e:
+        if isinstance(e, gr.Error):
+            raise
+        raise gr.Error(f"Error processing audio: {str(e)}")
+
+
+# Create the Gradio interface
+with gr.Blocks(title="Transcript Enhancer") as demo:
+    gr.Markdown("""
+    # 🎙️ Audio Transcript Enhancer
+    
+    Upload an audio file to get both an automated transcript and an enhanced version using AI.
+    
+    1. The original transcript is generated using AssemblyAI with speaker detection
+    2. The enhanced version uses Google's Gemini to improve clarity and readability
+    """)
+    
+    with gr.Row():
+        audio_input = gr.File(
+            label="Upload Audio File", 
+            type="binary", 
+            file_count="single", 
+            file_types=["audio"]
+        )
+    
+    with gr.Row():
+        transcribe_btn = gr.Button("📝 Transcribe & Enhance")
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Original Transcript")
+            original_output = gr.Markdown()
+        
+        with gr.Column():
+            gr.Markdown("### Enhanced Transcript")
+            enhanced_output = gr.Markdown()
+    
+    transcribe_btn.click(
+        fn=process_audio,
+        inputs=[audio_input],
+        outputs=[original_output, enhanced_output]
+    )
+
+# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(max_file_size=5 * gr.FileSize.GB)  # Backend limit