Blame for app.py

import spaces
import torch
import gradio as gr
import yt_dlp as youtube_dl
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
from transformers import WhisperProcessor
from pyannote.audio import Pipeline as PyannotePipeline
from datetime import datetime
import tempfile
import numpy as np
from itertools import groupby
import os
import re
import time  # Added missing import


os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:307'

try:
    diarization_pipeline = PyannotePipeline.from_pretrained(
        "pyannote/speaker-diarization-3.1",
        use_auth_token=os.environ["HF_TOKEN"]
    )
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    diarization_pipeline.to(device)
except Exception as e:
    print(f"Error initializing diarization pipeline: {e}")
    diarization_pipeline = None

# Updated Whisper model
MODEL_NAME = "openai/whisper-medium"
FILE_LIMIT_MB = 1000
YT_LENGTH_LIMIT_S = 3600

device = 0 if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    device=device,
    model_kwargs={"low_cpu_mem_usage": True},
)


def associate_speakers_with_timestamps(transcription_result, diarization):
    word_segments = transcription_result['chunks']
    diarization_segments = list(diarization.itertracks(yield_label=True))
    speaker_transcription = []
    current_segment_index = 0
    previous_speaker = None

    for word in word_segments:
        word_start, word_end = word['timestamp']

        while current_segment_index < len(diarization_segments) - 1 and diarization_segments[current_segment_index][0].end <= word_start:
            current_segment_index += 1

        current_segment, _, current_speaker = diarization_segments[current_segment_index]

        # Check for overlap with the previous segment
        if current_segment_index > 0:
            previous_segment, _, previous_speaker = diarization_segments[current_segment_index - 1]
            if previous_segment.end > word_start and previous_speaker != current_speaker:
                word_text = f"[{word['text']}]"
            else:
                word_text = word['text']
        else:
            word_text = word['text']

        if not speaker_transcription or speaker_transcription[-1][0] != current_speaker:
            speaker_transcription.append((current_speaker, word_text))
        else:
            speaker_transcription[-1] = (current_speaker, speaker_transcription[-1][1] + " " + word_text)

    return speaker_transcription
    
def simplify_diarization_output(speaker_transcription):
    simplified = []
    for speaker, text in speaker_transcription:
        simplified.append(f"{speaker}: {text}")
    return "\n".join(simplified)

def parse_simplified_diarization(simplified_text):
    pattern = r"(SPEAKER_\d+):\s*(.*)"
    matches = re.findall(pattern, simplified_text, re.MULTILINE)
    return [(speaker, text.strip()) for speaker, text in matches]

def count_unique_speakers(speaker_transcription):
    if isinstance(speaker_transcription, str):
        speaker_transcription = parse_simplified_diarization(speaker_transcription)
    return len(set(speaker for speaker, _ in speaker_transcription))

def process_transcription(*args):
    generator = transcribe_and_diarize(*args)
    for progress_message, raw_text, speaker_transcription in generator:
        pass  # Consume the generator until the end
    simplified_diarization = simplify_diarization_output(speaker_transcription)
    num_speakers = count_unique_speakers(speaker_transcription)
    speaker_names = [[f"SPEAKER_{i:02d}", f"SPEAKER_{i:02d}"] for i in range(num_speakers)]
    return progress_message, raw_text, simplified_diarization, num_speakers, speaker_names

def process_yt_transcription(*args):
    html_embed, raw_text, speaker_transcription = yt_transcribe(*args)
    simplified_diarization = simplify_diarization_output(speaker_transcription)
    num_speakers = count_unique_speakers(speaker_transcription)
    speaker_names = [[f"SPEAKER_{i:02d}", f"SPEAKER_{i:02d}"] for i in range(num_speakers)]
    return html_embed, raw_text, simplified_diarization, num_speakers, speaker_names
    

# New functions for progress indicator
def create_progress_indicator():
    return gr.State({"stage": 0, "message": "Waiting to start..."})

def update_progress(progress_state, stage, message):
    progress_state["stage"] = stage
    progress_state["message"] = message
    return progress_state

def display_progress(progress_state):
    stages = [
        "Loading file",
        "Preparing audio",
        "Transcribing",
        "Diarization (speaker identification)",
        "Finalizing results"
    ]
    progress = (progress_state["stage"] / len(stages)) * 100
    return gr.HTML(f"""
        <style>
            @keyframes pulse {{
                0% {{ box-shadow: 0 0 0 0 rgba(66, 133, 244, 0.7); }}
                70% {{ box-shadow: 0 0 0 10px rgba(66, 133, 244, 0); }}
                100% {{ box-shadow: 0 0 0 0 rgba(66, 133, 244, 0); }}
            }}
        </style>
        <div style="margin-bottom: 20px; background-color: #f0f0f0; padding: 15px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
            <h4 style="margin-bottom: 10px; color: #333; font-weight: bold;">Transcription Progress</h4>
            <div style="background-color: #e0e0e0; height: 24px; border-radius: 12px; overflow: hidden; position: relative;">
                <div style="background-color: #4285F4; width: {progress}%; height: 100%; border-radius: 12px; transition: width 0.5s ease-in-out;"></div>
                <div style="position: absolute; top: 0; left: 0; right: 0; bottom: 0; display: flex; align-items: center; justify-content: center; color: #333; font-weight: bold;">
                    {progress:.0f}%
                </div>
            </div>
            <p style="margin-top: 10px; color: #666; font-style: italic;">{progress_state["message"]}</p>
            <div style="width: 10px; height: 10px; background-color: #4285F4; border-radius: 50%; margin-top: 10px; animation: pulse 2s infinite;"></div>
        </div>
    """)
    
@spaces.GPU(duration=120)
def transcribe_and_diarize(file_path, task, progress=gr.Progress()):
    progress(0, desc="Initializing...")
    yield "Loading file...", None, None

    progress(0.2, desc="Preparing audio...")
    yield "Preparing audio...", None, None

    progress(0.4, desc="Give me a few minutes to decipher the voices and draft the audio 🤓 ✍️ ...")
    transcription_result = pipe(file_path, generate_kwargs={"task": task, "language": "fr"}, return_timestamps="word")
    yield "Transcribing...", None, None

    progress(0.6, desc="Done 😮‍💨! I'm working on merging everything, one moment, I'm almost there...")
    if diarization_pipeline:
        diarization = diarization_pipeline(file_path)
        speaker_transcription = associate_speakers_with_timestamps(transcription_result, diarization)
    else:
        speaker_transcription = [(None, transcription_result['text'])]
    yield "Diarizing...", None, None

    progress(0.8, desc="Finalizing results...")
    yield "Here it is!", transcription_result['text'], speaker_transcription

    progress(1.0, desc="Complete!")
    return "Transcription complete!", transcription_result['text'], speaker_transcription

def format_to_markdown(transcription_text, speaker_transcription, audio_duration=None, location=None, speaker_age=None, context=None, num_speakers=None, speaker_names=None):
    print("Starting format_to_markdown")
    print("Received speaker_names:", speaker_names)
    print("Received num_speakers:", num_speakers)
    
    metadata = {
        "Processing Date": datetime.now().strftime('%m/%d/%Y %H:%M'),
        "Audio Duration": f"{audio_duration}" if audio_duration else "[to be filled]",
        "Location": location if location else "[not specified]",
        "Number of speakers": num_speakers if num_speakers else "[not specified]",
        "Speaker Age": f"{speaker_age} years" if speaker_age else "[not specified]",
        "Context": context if context else "[not specified]"
    }
    
    # Create a dictionary for custom names
    custom_names = {}
    if speaker_names is not None and isinstance(speaker_names, list):
        for row in speaker_names:
            if isinstance(row, dict):
                speaker_id = row.get("Speaker ID")
                custom_name = row.get("Custom Name")
                if speaker_id and custom_name:
                    custom_names[speaker_id] = custom_name
            elif isinstance(row, (list, tuple)) and len(row) == 2:
                custom_names[row[0]] = row[1]
    elif speaker_names is not None:
        # Handling for cases where speaker_names is a DataFrame
        for index, row in speaker_names.iterrows():
            speaker_id = row.get("Speaker ID")
            custom_name = row.get("Custom Name")
            if speaker_id and custom_name:
                custom_names[speaker_id] = custom_name
    
    print("Created custom_names:", custom_names)
    
    metadata_text = "\n".join([f"- **{key}** : '{value}'" for key, value in metadata.items()])
    
    try:
        if isinstance(speaker_transcription, str):
            speaker_transcription = parse_simplified_diarization(speaker_transcription)
        
        print("speaker_transcription after parsing:", speaker_transcription)
        
        formatted_transcription = []
        for speaker, text in speaker_transcription:
            custom_name = custom_names.get(speaker, speaker)
            print(f"Using name: {custom_name} for speaker: {speaker}")
            formatted_transcription.append(f"**{custom_name}**: {text}")
        transcription_text = "\n\n".join(formatted_transcription)
    except Exception as e:
        print(f"Error formatting speaker transcription: {e}")
        transcription_text = "Error formatting speaker transcription. Using raw transcription instead.\n\n" + transcription_text

    formatted_output = f"""
# Formatted Transcription

## Metadata
{metadata_text}

## Transcription
{transcription_text}
"""
    print("End of format_to_markdown")
    return formatted_output

    
def _return_yt_html_embed(yt_url):
    video_id = yt_url.split("?v=")[-1]
    HTML_str = (
        f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
        " </center>"
    )
    return HTML_str

def download_yt_audio(yt_url, filename):
    info_loader = youtube_dl.YoutubeDL()
    
    try:
        info = info_loader.extract_info(yt_url, download=False)
    except youtube_dl.utils.DownloadError as err:
        raise gr.Error(str(err))
    
    file_length = info["duration"]
    
    if file_length > YT_LENGTH_LIMIT_S:
        yt_length_limit_hms = time.strftime("%H:%M:%S", time.gmtime(YT_LENGTH_LIMIT_S))
        file_length_hms = time.strftime("%H:%M:%S", time.gmtime(file_length))
        raise gr.Error(f"Maximum YouTube duration is {yt_length_limit_hms}, the YouTube video is {file_length_hms} long.")
    
    ydl_opts = {"outtmpl": filename, "format": "bestaudio/best"}
    
    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        try:
            ydl.download([yt_url])
        except youtube_dl.utils.ExtractorError as err:
            raise gr.Error(str(err))

@spaces.GPU(duration=120)
def yt_transcribe(yt_url, task):
    html_embed_str = _return_yt_html_embed(yt_url)

    with tempfile.TemporaryDirectory() as tmpdirname:
        filepath = os.path.join(tmpdirname, "video.mp4")
        download_yt_audio(yt_url, filepath)
        with open(filepath, "rb") as f:
            inputs = f.read()

    inputs = ffmpeg_read(inputs, pipe.feature_extractor.sampling_rate)
    inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}

    transcription_result = pipe(inputs, generate_kwargs={"task": task}, return_timestamps=True)
    transcription_text = transcription_result['text']

    if diarization_pipeline:
        diarization = diarization_pipeline(filepath)
        speaker_transcription = associate_speakers_with_timestamps(transcription_result, diarization)
    else:
        speaker_transcription = [(None, transcription_text)]

    return html_embed_str, transcription_text, speaker_transcription

def create_info_box(title, content):
    return gr.Markdown(f"### {title}\n{content}")

css_file_path = os.path.join(os.path.dirname(__file__), "style.css")
with open(css_file_path, "r") as f:
    custom_css = f.read()

# Gradio 6.x: theme and css moved to launch()
demo = gr.Blocks(
    title="Scribe - Intelligent Audio Transcription Assistant 🎙️📝",
)


with demo:
    gr.Markdown("# 🎙️ **Scribe**: Intelligent Audio Transcription Assistant 📝 ")
    gr.Markdown('# ⚠️ This is a public demo. For confidential use, please click the button below 👇')
    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
    gr.HTML(
        """
        <div class="logo">
            <img src="https://image.noelshack.com/fichiers/2024/33/4/1723713257-dbe58773-0638-445b-a88c-3fc1f2002408.jpg" alt="Scribe Logo">
        </div>
        """
    )
    gr.Markdown("## **Welcome to Scribe, a secure audio transcription solution. Efficiently transform your audio files, live recordings, or YouTube videos into accurate text.**")

    gr.Markdown("""
    ### 🔍 **How the Model Works**:
    Scribe uses a two-step approach to transform audio into structured text:

    1. **Transcription with Whisper Medium**:
       - Speech recognition model developed by OpenAI
       - Uses an encoder-decoder neural network with attention
       - Capable of handling various accents and background noise
       - Optimized for a balance between accuracy and speed

    2. **Diarization with pyannote/speaker-diarization-3.1**:
       - Identifies and segments different speakers in the audio
       - Uses deep learning techniques for vocal feature extraction
       - Applies a clustering algorithm to group segments by speaker

    ### 💡 **Tips for Better Results**
    - Use high-quality recordings with minimal background noise.
    - For long recordings, it is recommended to segment your audio.
    - Always verify the transcription, especially for technical terms or proper names.
    - Use external microphones for live recordings if possible.

    ### ⚙️ Technical Specifications:
    - Transcription model: Whisper Medium
    - Diarization pipeline: pyannote/speaker-diarization-3.1
    - File size limit: _(We currently have no precise limit. However, **we recommend not exceeding 5 minutes.** )_
    - Maximum duration for YouTube videos: _(We currently have no precise limit. However, for optimal use, audio should not exceed 30 minutes. )_
    - Supported audio formats: MP3, WAV, M4A, and more
    """)
    with gr.Accordion("🔐 Data Security and Pipelines", open=False):
        gr.Markdown("""
    #### What is a pipeline?
    In the context of machine learning, a pipeline is a series of data processing steps, from raw input to final output. In Scribe, we use two main pipelines:

    1. **Transcription Pipeline**: Based on the Whisper Medium model, it converts audio to text.
    2. **Diarization Pipeline**: Identifies different speakers in the audio.

    #### How do our pipelines work?
    1. **Local Loading**: Models are loaded locally on your machine or server.
    2. **In-Situ Processing**: All data is processed locally, without being sent to external servers.
    3. **Volatile Memory**: Data is temporarily stored in RAM and cleared after use.

    #### Security and Privacy
    - **No External Transmission**: Your audio and text data remain on your local system.
    - **Isolation**: Each user session is isolated from others.
    - **Automatic Cleanup**: Temporary files are deleted after each use.
    - **Secure Updates**: Models are securely updated via Hugging Face.

    #### Additional Security Measures
    - We use secure authentication tokens to access the models.
    - YouTube files are downloaded and processed locally, without permanent storage.
    - No user data is retained after the session ends.

    By using Scribe, you benefit from highly secure and privacy-respecting data processing, while leveraging the power of cutting-edge AI models.
    """)
    
    with gr.Tabs():
        with gr.Tab("Audio File 📁"):
            gr.Markdown("### 📂 Audio File Transcription")
            audio_input = gr.Audio(type="filepath", label="Upload your audio file")
            task_input = gr.Radio(["transcribe", "translate"], label="Choose task", value="transcribe")
            transcribe_button = gr.Button("🚀 Start Transcription", elem_classes="button-primary")
            
            progress_display = gr.Markdown(label="Progress Status")
            
            with gr.Accordion("Results 📊", open=True):
                raw_output = gr.Textbox(label="📝 Raw Transcription", info="Text generated by the model. Editable if necessary.")
                speaker_output = gr.Textbox(label="👥 Diarization (simplified format)", info="Speaker identification. Format: 'SPEAKER_XX: text'")
            with gr.Accordion("Metadata (optional) 📌", open=False):
                audio_duration = gr.Textbox(label="⏱️ Audio Duration (mm:ss)")
                location = gr.Textbox(label="📍 Recording Location")
                num_speakers = gr.Number(label="Number of speakers", interactive=False)
                # Gradio 6.x: col_count replaced with column_count and column_limits
                speaker_names = gr.Dataframe(
                    headers=["Speaker ID", "Custom Name"],
                    datatype=["str", "str"],
                    column_count=2,
                    column_limits=(2, 2),
                    label="Speaker Names"
                )
                speaker_age = gr.Number(label="👤 Main Speaker Age")
                context = gr.Textbox(label="📝 Recording Context")
            
            format_button = gr.Button("✨ Generate Formatted Transcription", elem_classes="button-secondary")
            formatted_output = gr.Markdown(label="📄 Formatted Transcription:")


        with gr.Tab("Microphone 🎤"):
            gr.Markdown("### 🗣️ Live Recording and Transcription")
            mic_input = gr.Audio(type="filepath", label="Record your voice")
            mic_task_input = gr.Radio(["transcribe", "translate"], label="Choose task", value="transcribe")
            mic_transcribe_button = gr.Button("🚀 Transcribe Recording", elem_classes="button-primary")
            
            mic_progress_display = gr.Markdown(label="Progress Status")
            
            with gr.Accordion("Results 📊", open=True):
                mic_raw_output = gr.Textbox(label="📝 Raw Transcription", info="Text generated by the model. Editable if necessary.")
                mic_speaker_output = gr.Textbox(label="👥 Diarization (simplified format)", info="Speaker identification. Format: 'SPEAKER_XX: text'")
            with gr.Accordion("Metadata (optional) 📌", open=False):
                mic_audio_duration = gr.Textbox(label="⏱️ Recording Duration (mm:ss)")
                mic_location = gr.Textbox(label="📍 Recording Location")
                mic_num_speakers = gr.Number(label="Number of speakers", interactive=False)
                mic_speaker_names = gr.Dataframe(
                    headers=["Speaker ID", "Custom Name"],
                    datatype=["str", "str"],
                    column_count=2,
                    column_limits=(2, 2),
                    label="Speaker Names"
                )
                mic_speaker_age = gr.Number(label="👤 Main Speaker Age")
                mic_context = gr.Textbox(label="📝 Recording Context")
            
            mic_format_button = gr.Button("✨ Generate Formatted Transcription", elem_classes="button-secondary")
            mic_formatted_output = gr.Markdown(label="📄 Formatted Transcription:")
            
        with gr.Tab("YouTube 🎥"):
            gr.Markdown("### 🌐 Transcription from YouTube Videos")
            yt_input = gr.Textbox(lines=1, placeholder="Paste a YouTube video URL here", label="🔗 YouTube URL")
            yt_task_input = gr.Radio(["transcribe", "translate"], label="Choose task", value="transcribe")
            yt_transcribe_button = gr.Button("🚀 Transcribe Video", elem_classes="button-primary")
            
            yt_progress_display = gr.Markdown(label="Progress Status")
            
            yt_html_output = gr.HTML(label="▶️ Video Preview")
            
            with gr.Accordion("Results 📊", open=True):
                yt_raw_output = gr.Textbox(label="📝 Raw Transcription", info="Text generated by the model. Editable if necessary.")
                yt_speaker_output = gr.Textbox(label="👥 Diarization (simplified format)", info="Speaker identification. Format: 'SPEAKER_XX: text'")
            with gr.Accordion("Metadata (optional) 📌", open=False):
                yt_audio_duration = gr.Textbox(label="⏱️ Video Duration (mm:ss)")
                yt_channel = gr.Textbox(label="📺 YouTube Channel Name")
                yt_num_speakers = gr.Number(label="Number of speakers", interactive=False)
                yt_speaker_names = gr.Dataframe(
                    headers=["Speaker ID", "Custom Name"],
                    datatype=["str", "str"],
                    column_count=2,
                    column_limits=(2, 2),
                    label="Speaker Names"
                )
                yt_publish_date = gr.Textbox(label="📅 Publication Date")
                yt_context = gr.Textbox(label="📝 Video Context")
            
            yt_format_button = gr.Button("✨ Generate Formatted Transcription", elem_classes="button-secondary")
            yt_formatted_output = gr.Markdown(label="📄 Formatted Transcription:")


    gr.Markdown("""### 🛠️ Capabilities:
    - Multilingual transcription with automatic language detection
    - Translation to English (for non-English content)
    - Accurate identification of speaker changes
    - Processing of audio files, live recordings, and YouTube videos
    - Handling of various audio formats and recording qualities
    """)
        
    with gr.Accordion("❓ README:", open=False):
        gr.Markdown("""
    - Creator: Woziii
    - Models:
        - [Whisper-medium](https://huggingface.co/openai/whisper-medium) : Model size - 764M params - Tensor type F32 -
        - [speaker-diarization-3.1](https://huggingface.co/pyannote/speaker-diarization-3.1) : Model size - Unknown - Tensor type F32 -
    - Version: V.2.1.3-Beta
    - Languages: EN, FR
    - Copyright: cc-by-nc-4.0
    - [Learn more](https://huggingface.co/spaces/Woziii/scribe/blob/main/README.md)
    """)

    # Button connections to appropriate functions
    transcribe_button.click(
        process_transcription,
        inputs=[audio_input, task_input],
        outputs=[progress_display, raw_output, speaker_output, num_speakers, speaker_names]
    )
    
    format_button.click(
        format_to_markdown,
        inputs=[raw_output, speaker_output, audio_duration, location, speaker_age, context, num_speakers, speaker_names],
        outputs=formatted_output
    )
    
    mic_transcribe_button.click(
        process_transcription,
        inputs=[mic_input, mic_task_input],
        outputs=[mic_progress_display, mic_raw_output, mic_speaker_output, mic_num_speakers, mic_speaker_names]
    )
    
    mic_format_button.click(
        format_to_markdown,
        inputs=[mic_raw_output, mic_speaker_output, mic_audio_duration, mic_location, mic_speaker_age, mic_context, mic_num_speakers, mic_speaker_names],
        outputs=mic_formatted_output
    )
    
    yt_transcribe_button.click(
        process_yt_transcription,
        inputs=[yt_input, yt_task_input],
        outputs=[yt_html_output, yt_raw_output, yt_speaker_output, yt_num_speakers, yt_speaker_names]
    )
    
    yt_format_button.click(
        format_to_markdown,
        inputs=[yt_raw_output, yt_speaker_output, yt_audio_duration, yt_channel, yt_publish_date, yt_context, yt_num_speakers, yt_speaker_names],
        outputs=yt_formatted_output
    )
        

if __name__ == "__main__":
    # Gradio 6.x: theme and css parameters moved to launch()
    demo.queue().launch(
        theme=gr.themes.Default(),
        css=custom_css
    )