Update app.py

app.py

@@ -3,127 +3,146 @@ import torch
 from transformers import pipeline
 from docx import Document
 import datetime
-import time
 from pydub import AudioSegment
 import os
 
 # --- Model setup ---
-MODEL_ID = "openai/whisper-tiny"  # Tiny Whisper model
+MODEL_ID = "openai/whisper-tiny"
 FRENCH_LANGUAGE = "fr"
-CHUNK_LENGTH_SEC = 600  # 10 minutes
+CHUNK_LENGTH_SEC = 600
 
-# Cache for the pipeline
 model_cache = {}
 
-def get_model_pipeline(progress=None):
+def get_model_pipeline():
     if "tiny_fr" not in model_cache:
-        if progress: progress(0, desc="🚀 Loading model...")
+        device = 0 if torch.cuda.is_available() else -1
         model_cache["tiny_fr"] = pipeline(
             "automatic-speech-recognition",
             model=MODEL_ID,
-            device=-1,  # Force CPU
+            device=device,
+            torch_dtype=torch.float16 if device >= 0 else torch.float32,
             chunk_length_s=30
         )
-        if progress: progress(1.0, desc="✅ Model loaded!")
     return model_cache["tiny_fr"]
 
-# --- Audio slicing ---
-def slice_audio(audio_path, chunk_length_sec=CHUNK_LENGTH_SEC):
+# --- Audio slicing (optional for GPU) ---
+def slice_audio(audio_path):
     audio = AudioSegment.from_file(audio_path)
-    chunks = []
-    for i in range(0, len(audio), chunk_length_sec * 1000):  # pydub uses milliseconds
-        chunks.append(audio[i:i + chunk_length_sec * 1000])
-    return chunks
+    slices = []
+    for i in range(0, len(audio), CHUNK_LENGTH_SEC*1000):
+        slices.append(audio[i:i+CHUNK_LENGTH_SEC*1000])
+    return slices
 
-def export_chunks(chunks, base_name="slice"):
+def export_chunks(chunks):
     paths = []
     for idx, chunk in enumerate(chunks):
-        file_path = f"{base_name}_{idx+1}.wav"
-        chunk.export(file_path, format="wav")
-        paths.append(file_path)
+        path = f"slice_{idx+1}.wav"
+        chunk.export(path, format="wav")
+        paths.append(path)
     return paths
 
-# --- Output functions ---
-def create_vtt(segments, file_path):
-    with open(file_path, "w", encoding="utf-8") as f:
-        f.write("WEBVTT\n\n")
-        for i, s in enumerate(segments):
-            start = str(datetime.timedelta(seconds=int(s.get("start", 0))))
-            end = str(datetime.timedelta(seconds=int(s.get("end", 0))))
-            f.write(f"{i+1}\n{start} --> {end}\n{s.get('text','').strip()}\n\n")
-
-def create_docx(segments, file_path, with_timestamps=True):
+# --- DOCX helpers ---
+def create_docx_from_segments(segments, file_path, with_timestamps=True):
     doc = Document()
-    doc.add_heading("Transcription", 0)
-    if with_timestamps:
-        for s in segments:
-            start = str(datetime.timedelta(seconds=int(s.get("start", 0))))
-            end = str(datetime.timedelta(seconds=int(s.get("end", 0))))
-            doc.add_paragraph(f"[{start} - {end}] {s.get('text','').strip()}")
-    else:
-        doc.add_paragraph(" ".join([s.get("text",'').strip() for s in segments]))
+    doc.add_heading("Full Transcription", 0)
+    for seg in segments:
+        text = seg.get("text", "").strip()
+        start = str(datetime.timedelta(seconds=int(seg.get("start", 0))))
+        end = str(datetime.timedelta(seconds=int(seg.get("end", 0))))
+        if with_timestamps:
+            doc.add_paragraph(f"[{start} - {end}] {text}")
+        else:
+            doc.add_paragraph(text)
     doc.save(file_path)
 
-# --- Transcription function for slices ---
-def transcribe_slice(audio_file, vtt=True, docx_ts=False, docx_no_ts=True, progress=gr.Progress()):
-    if audio_file is None:
-        return None, None, "Please upload an audio file."
-
-    start_time = time.time()
-    pipe = get_model_pipeline(progress)
-
-    # Split audio into 10-minute slices
-    chunks = slice_audio(audio_file)
-    chunk_files = export_chunks(chunks, "slice")
+def create_docx_from_text(text_list, file_path):
+    doc = Document()
+    doc.add_heading("Patchwork Summary", 0)
+    for text in text_list:
+        doc.add_paragraph(text)
+    doc.save(file_path)
 
+# --- Step 1: Slice audio ---
+def generate_slice_options(audio_file):
+    slices = slice_audio(audio_file)
+    chunk_files = export_chunks(slices)
+    options = []
+    for idx, chunk in enumerate(slices):
+        start_min = idx*10
+        end_min = start_min + min(10, len(chunk)//1000//60)
+        options.append(f"Slice {idx+1}: {start_min}-{end_min} min")
+    return options, chunk_files
+
+# --- Step 2: Transcribe selected slices ---
+def transcribe_selected_slices(selected_slices, chunk_files):
+    if not selected_slices:
+        return "", "Please select at least one slice.", [], []
+
+    pipe = get_model_pipeline()
     full_text = ""
-    outputs = {}
     all_segments = []
 
-    for idx, chunk_file in enumerate(chunk_files):
-        progress((idx+1)/len(chunk_files), desc=f"🎤 Transcribing slice {idx+1}/{len(chunk_files)}...")
-        result = pipe(chunk_file, return_timestamps=True, generate_kwargs={"language": FRENCH_LANGUAGE})
+    for sel in selected_slices:
+        idx = int(sel.split()[1].replace(":", "")) - 1
+        result = pipe(chunk_files[idx], return_timestamps=True, generate_kwargs={"language": FRENCH_LANGUAGE})
         segments = result.get("chunks", [])
         all_segments.extend(segments)
         full_text += result["text"] + "\n"
 
-    progress(0.85, desc="📝 Generating files...")
+    sentences = [s.get("text","").strip() for s in all_segments if s.get("text","").strip()]
+    return full_text, "✅ Transcription complete!", sentences, all_segments
+
+# --- Step 3: Generate DOCX ---
+def generate_docs(all_segments, selected_sentences):
+    full_docx_path = "full_transcription.docx"
+    create_docx_from_segments(all_segments, full_docx_path, with_timestamps=True)
 
-    if vtt:
-        outputs["VTT"] = "transcription.vtt"
-        create_vtt(all_segments, outputs["VTT"])
-    if docx_ts:
-        outputs["DOCX_TS"] = "transcription_with_timestamps.docx"
-        create_docx(all_segments, outputs["DOCX_TS"], True)
-    if docx_no_ts:
-        outputs["DOCX_NO_TS"] = "transcription_without_timestamps.docx"
-        create_docx(all_segments, outputs["DOCX_NO_TS"], False)
+    patchwork_docx_path = None
+    if selected_sentences:
+        patchwork_docx_path = "patchwork_summary.docx"
+        create_docx_from_text(selected_sentences, patchwork_docx_path)
 
-    total_time = time.time() - start_time
-    return full_text, list(outputs.values()), f"✅ Done! Took {total_time:.2f} seconds."
+    files = [full_docx_path]
+    if patchwork_docx_path:
+        files.append(patchwork_docx_path)
+
+    return files, "✅ DOCX files generated!"
 
 # --- Gradio UI ---
-with gr.Blocks(title="Tiny French Whisper CPU with 10-min Slices") as demo:
-    gr.Markdown("# 🎙️ Tiny Whisper - French CPU Transcription")
-    gr.Markdown("Upload your audio file and get a transcription sliced into 10-minute segments, optimized for French (CPU).")
-
-    with gr.Row():
-        audio_input = gr.Audio(label="Audio File", type="filepath")  # Fixed for latest Gradio
-        with gr.Column():
-            vtt_cb = gr.Checkbox(label="VTT", value=True)
-            docx_ts_cb = gr.Checkbox(label="DOCX (with timestamps)", value=False)
-            docx_no_ts_cb = gr.Checkbox(label="DOCX (without timestamps)", value=True)
-            transcribe_btn = gr.Button("Transcribe", variant="primary")
-            status_out = gr.Textbox(label="Status", interactive=False)
-
-    transcription_out = gr.Textbox(label="Full Transcription", lines=10)
-    files_out = gr.Files(label="Download Transcripts")
-
-    transcribe_btn.click(
-        fn=transcribe_slice,
-        inputs=[audio_input, vtt_cb, docx_ts_cb, docx_no_ts_cb],
-        outputs=[transcription_out, files_out, status_out]
-    )
+with gr.Blocks(title="Tiny Whisper GPU French Transcription") as demo:
+    gr.Markdown("## 🎙️ GPU-Optimized Tiny Whisper - French Transcription with Summary")
+
+    with gr.Tab("Step 1: Upload & Slice Audio"):
+        audio_input = gr.Audio(label="Upload Audio", type="filepath")
+        generate_slices_btn = gr.Button("Generate Slices")
+        slices_output = gr.CheckboxGroup(label="Select slices to transcribe")
+        chunk_files_state = gr.State()
+        generate_slices_btn.click(
+            generate_slice_options,
+            inputs=[audio_input],
+            outputs=[slices_output, chunk_files_state]
+        )
+
+    with gr.Tab("Step 2: Transcribe Selected Slices"):
+        transcribe_btn = gr.Button("Transcribe Selected Slices")
+        transcription_output = gr.Textbox(label="Full Transcription", lines=10)
+        sentence_selector = gr.CheckboxGroup(label="Select sentences for patchwork")
+        status_transcription = gr.Textbox(label="Status", interactive=False)
+        transcribe_btn.click(
+            transcribe_selected_slices,
+            inputs=[slices_output, chunk_files_state],
+            outputs=[transcription_output, status_transcription, sentence_selector, gr.State()]
+        )
+
+    with gr.Tab("Step 3: Generate DOCX"):
+        generate_docs_btn = gr.Button("Generate DOCX Files")
+        docs_output = gr.Files(label="Download DOCX files")
+        docs_status = gr.Textbox(label="Status", interactive=False)
+        generate_docs_btn.click(
+            generate_docs,
+            inputs=[gr.State(), sentence_selector],
+            outputs=[docs_output, docs_status]
+        )
 
 if __name__ == "__main__":
     demo.launch()