Update app.py

app.py

@@ -1,6 +1,6 @@
 # app.py
 # Whisper transcription app - HYBRID conversion (pydub + small ffmpeg fallback)
-# Clean, single-version file for Hugging Face Spaces.
+# Multi-tab UI, zip extraction + selectable files, memory management
 
 import os
 import sys
@@ -13,12 +13,12 @@ import threading
 import re
 from difflib import get_close_matches
 
-# Force unbuffered output so container logs show prints immediately
+# Force unbuffered output
 os.environ["PYTHONUNBUFFERED"] = "1"
 
 print("DEBUG: app.py bootstrap starting", flush=True)
 
-# Third-party imports (must be installed in the environment)
+# Third-party imports
 try:
     from docx import Document
     import whisper
@@ -35,7 +35,7 @@ print("DEBUG: imports OK", flush=True)
 # ---------- Config ----------
 MEMORY_FILE = "memory.json"
 MEMORY_LOCK = threading.Lock()
-MIN_WAV_SIZE = 1024  # bytes
+MIN_WAV_SIZE = 1024
 FFMPEG_CANDIDATES = [
     ("s16le", 16000, 1),
     ("s16le", 44100, 2),
@@ -77,13 +77,9 @@ def save_memory(mem):
 
 
 memory = load_memory()
-print(
-    "DEBUG: memory loaded (words=%d phrases=%d)"
-    % (len(memory.get("words", {})), len(memory.get("phrases", {}))),
-    flush=True,
-)
+print("DEBUG: memory loaded (words=%d phrases=%d)" % (len(memory.get("words", {})), len(memory.get("phrases", {}))), flush=True)
 
-# ---------- Postprocessing ----------
+# ---------- Postprocessing (same as before) ----------
 MEDICAL_ABBREVIATIONS = {
     "pt": "patient",
     "dx": "diagnosis",
@@ -159,14 +155,12 @@ def postprocess_transcript(text, format_soap=False):
             if kw in t.lower():
                 assessment = "Assessment: " + subj
                 break
-        soap = (
-            f"S: {subj}\nO: {obj}\nA: {assessment}\nP: Plan: follow up as indicated."
-        )
+        soap = f"S: {subj}\nO: {obj}\nA: {assessment}\nP: Plan: follow up as indicated."
         return soap
     return t
 
 
-# ---------- Memory utilities ----------
+# ---------- Memory utilities (same as before) ----------
 def extract_words_and_phrases(text):
     words = re.findall(r"[A-Za-z0-9\-']+", text)
     sentences = [s.strip() for s in re.split(r"(?<=[.?!])\s+", text) if s.strip()]
@@ -234,7 +228,7 @@ def memory_correct_text(text, min_ratio=0.85):
     return corrected
 
 
-# ---------- Memory management UI helpers ----------
+# ---------- Memory management UI helpers (same as before) ----------
 def import_memory_file(uploaded):
     global memory
     if not uploaded:
@@ -319,7 +313,7 @@ def clear_memory():
     return "Memory cleared."
 
 
-def view_memory(limit=2000):
+def view_memory(limit=4000):
     w = memory.get("words", {})
     p = memory.get("phrases", {})
     out_lines = []
@@ -461,136 +455,105 @@ def convert_to_wav_if_needed(input_path):
 MODEL_CACHE = {}
 
 
-def get_whisper_model(name):
+def get_whisper_model(name, device=None):
     if name not in MODEL_CACHE:
         print(f"DEBUG: loading whisper model '{name}'", flush=True)
-        MODEL_CACHE[name] = whisper.load_model(name)
+        if device:
+            MODEL_CACHE[name] = whisper.load_model(name, device=device)
+        else:
+            MODEL_CACHE[name] = whisper.load_model(name)
     return MODEL_CACHE[name]
 
 
-# ---------- Main transcription generator ----------
+# ---------- ZIP extraction + selection helpers ----------
+def extract_zip_list(zip_file, zip_password):
+    """
+    Extract zip to a temp dir and return (list_of_paths, diagnostics_text)
+    """
+    temp_extract_dir = os.path.join(tempfile.gettempdir(), "extracted_audio")
+    try:
+        if os.path.exists(temp_extract_dir):
+            # clear existing
+            try:
+                shutil.rmtree(temp_extract_dir)
+            except Exception:
+                pass
+        os.makedirs(temp_extract_dir, exist_ok=True)
+        extracted = []
+        logs = []
+        with pyzipper.ZipFile(zip_file, "r") as zf:
+            if zip_password:
+                try:
+                    zf.setpassword(zip_password.encode())
+                except Exception:
+                    logs.append("Warning: failed to set zip password (unexpected).")
+            exts = [".mp3", ".wav", ".aac", ".flac", ".ogg", ".m4a", ".dat", ".dct"]
+            for info in zf.infolist():
+                if info.is_dir():
+                    continue
+                _, ext = os.path.splitext(info.filename)
+                if ext.lower() in exts:
+                    try:
+                        zf.extract(info, path=temp_extract_dir)
+                    except RuntimeError as e:
+                        logs.append(f"Password required/incorrect for {info.filename}: {e}")
+                        continue
+                    except pyzipper.BadZipFile:
+                        logs.append(f"Bad zip entry: {info.filename}")
+                        continue
+                    except Exception as e:
+                        logs.append(f"Error extracting {info.filename}: {e}")
+                        continue
+                    p = os.path.normpath(os.path.join(temp_extract_dir, info.filename))
+                    if os.path.exists(p):
+                        extracted.append(p)
+                        logs.append(f"Extracted: {info.filename}")
+        if not extracted:
+            logs.append("No supported audio files found in zip.")
+            return [], "\n".join(logs)
+        # Return list and logs
+        return extracted, "\n".join(logs)
+    except Exception as e:
+        traceback.print_exc()
+        return [], f"Extraction failed: {e}"
+
+
+# ---------- Main transcription generator (updated to accept explicit 'selected_paths') ----------
 def transcribe_multiple(
-    audio_files,
+    selected_paths,
     model_name,
     advanced_options,
     merge_checkbox,
-    zip_file=None,
-    zip_password=None,
     enable_memory=False,
+    device=None,
 ):
+    """
+    Generator yields (log_text, transcripts_text, merged_file_path_or_None, percent_int)
+    selected_paths: list of absolute file paths to process
+    """
     log = []
     transcripts = []
     word_file_path = None
-    temp_extract_dir = os.path.join(tempfile.gettempdir(), "extracted_audio")
-    extracted_audio_paths = []
-
-    yield "", "", None, 0
-
-    if os.path.exists(temp_extract_dir):
-        try:
-            shutil.rmtree(temp_extract_dir)
-            log.append(f"Cleaned previous temp dir: {temp_extract_dir}")
-        except Exception:
-            pass
 
-    if zip_file:
-        log.append(f"Processing zip: {zip_file}")
-        yield "\n\n".join(log), "\n\n".join(transcripts), None, 2
-        try:
-            os.makedirs(temp_extract_dir, exist_ok=True)
-            with pyzipper.ZipFile(zip_file, "r") as zf:
-                if zip_password:
-                    try:
-                        zf.setpassword(zip_password.encode())
-                    except Exception:
-                        log.append("Failed to set zip password (unexpected).")
-
-                exts = [
-                    ".mp3",
-                    ".wav",
-                    ".aac",
-                    ".flac",
-                    ".ogg",
-                    ".m4a",
-                    ".dat",
-                    ".dct",
-                ]
-                count = 0
-                for info in zf.infolist():
-                    if info.is_dir():
-                        continue
-                    _, ext = os.path.splitext(info.filename)
-                    if ext.lower() in exts:
-                        try:
-                            zf.extract(info, path=temp_extract_dir)
-                        except RuntimeError as e:
-                            log.append(f"Password required or incorrect for {info.filename}: {e}")
-                            continue
-                        except pyzipper.BadZipFile:
-                            log.append(f"Bad zip entry: {info.filename}")
-                            continue
-                        except Exception as e:
-                            log.append(f"Error extracting {info.filename}: {e}")
-                            continue
-                        p = os.path.normpath(os.path.join(temp_extract_dir, info.filename))
-                        if os.path.exists(p):
-                            extracted_audio_paths.append(p)
-                            count += 1
-                            log.append(f"Extracted: {info.filename}")
-                if count == 0:
-                    log.append("No supported audio in zip.")
-                    try:
-                        shutil.rmtree(temp_extract_dir)
-                    except Exception:
-                        pass
-                    yield "\n\n".join(log), "\n\n".join(transcripts), None, 100
-                    return
-        except pyzipper.BadZipFile:
-            log.append("Invalid zip file.")
-            try:
-                shutil.rmtree(temp_extract_dir)
-            except Exception:
-                pass
-            yield "\n\n".join(log), "\n\n".join(transcripts), None, 100
-            return
-        except Exception as e:
-            log.append(f"Zip processing error: {e}")
-            try:
-                shutil.rmtree(temp_extract_dir)
-            except Exception:
-                pass
-            yield "\n\n".join(log), "\n\n".join(transcripts), None, 100
-            return
-
-    paths = []
-    if extracted_audio_paths:
-        paths.extend(extracted_audio_paths)
-    if audio_files:
-        if isinstance(audio_files, (list, tuple)):
-            for a in audio_files:
-                if a:
-                    paths.append(a)
-        elif isinstance(audio_files, str):
-            paths.append(audio_files)
-
-    if not paths:
+    if not selected_paths:
         log.append("No audio files provided.")
         yield "\n\n".join(log), "\n\n".join(transcripts), None, 100
         return
 
+    yield "", "", None, 0
+
+    # load model
     yield "\n\n".join(log), "\n\n".join(transcripts), None, 5
     try:
-        model = get_whisper_model(model_name)
+        model = get_whisper_model(model_name, device=device)
         log.append(f"Loaded Whisper model: {model_name}")
     except Exception as e:
         log.append(f"Failed to load model {model_name}: {e}")
         yield "\n\n".join(log), "\n\n".join(transcripts), None, 100
         return
 
-    total = len(paths)
-    idx = 0
-    for p in paths:
-        idx += 1
+    total = len(selected_paths)
+    for idx, p in enumerate(selected_paths, start=1):
         log.append(f"Processing file ({idx}/{total}): {p}")
         yield "\n\n".join(log), "\n\n".join(transcripts), None, int(5 + (idx - 1) * 80 / max(1, total))
 
@@ -659,188 +622,213 @@ def transcribe_multiple(
 
     yield "\n\n".join(log), "\n\n".join(transcripts), word_file_path, 100
 
-    try:
-        if os.path.exists(temp_extract_dir):
-            shutil.rmtree(temp_extract_dir)
-            log.append("Cleaned temporary extraction dir.")
-    except Exception:
-        pass
-
-
-# ----------------------- Gradio wrapper (streaming) -----------------------
-def run_transcription_wrapper(
-    files,
-    model_name,
-    merge,
-    zip_file,
-    zip_password,
-    use_default_zip_pass,
-    default_zip_password,
-    enable_memory,
-    advanced_options_state,
-):
-    try:
-        audio_input = files
 
-        zip_path = None
-        if zip_file:
+# ----------------------- Gradio callbacks & UI -----------------------
+def extract_zip_for_ui(zip_file, zip_password, use_default_zip_pass, default_zip_password):
+    """
+    Extract zip and return (checkbox_choices, logs)
+    """
+    if use_default_zip_pass and (not zip_password or zip_password.strip() == ""):
+        final_zip_password = default_zip_password
+    else:
+        final_zip_password = zip_password
+    if not zip_file:
+        return [], "No ZIP file provided."
+    # Normalize zip path
+    zip_path = None
+    if isinstance(zip_file, (str, os.PathLike)):
+        zip_path = str(zip_file)
+    elif hasattr(zip_file, "name"):
+        zip_path = zip_file.name
+    elif isinstance(zip_file, dict) and zip_file.get("name"):
+        zip_path = zip_file["name"]
+    else:
+        return [], "Unable to determine uploaded zip path."
+
+    extracted, logs = extract_zip_list(zip_path, final_zip_password)
+    # For the UI we show readable labels but the choices list will hold full paths
+    choices = extracted  # list of paths (strings)
+    return choices, logs or "Extraction completed."
+
+
+def run_transcription_ui(selected_files, file_input, model_name, merge, zip_selected_files, zip_file, zip_password, use_default_zip_pass, default_zip_password, enable_memory, device_choice):
+    """
+    Top-level UI handler invoked by the Transcribe button.
+    Priority:
+      1) zip_selected_files: explicit selection of extracted files (checkbox group)
+      2) selected_files from file_input (file input paths)
+      3) zip_file without explicit selection -> extract all then transcribe
+    This function returns a Gradio generator (yields) using transcribe_multiple.
+    """
+    # build final list of files to process
+    final_paths = []
+
+    # If the user selected extracted zip files (zip_selected_files is list of paths), use those
+    if zip_selected_files:
+        final_paths = zip_selected_files if isinstance(zip_selected_files, (list, tuple)) else [zip_selected_files]
+    else:
+        # if file_input provided (list of paths), use them
+        if file_input:
+            if isinstance(file_input, (list, tuple)):
+                for a in file_input:
+                    if a:
+                        # file_input uses type="filepath" so entries are paths
+                        final_paths.append(str(a))
+            elif isinstance(file_input, str):
+                final_paths.append(file_input)
+
+        # if nothing chosen and zip_file provided, auto-extract all and use them
+        if not final_paths and zip_file:
+            # reuse extract logic
+            if use_default_zip_pass and (not zip_password or zip_password.strip() == ""):
+                final_zip_password = default_zip_password
+            else:
+                final_zip_password = zip_password
+            zip_path = None
             if isinstance(zip_file, (str, os.PathLike)):
                 zip_path = str(zip_file)
             elif hasattr(zip_file, "name"):
                 zip_path = zip_file.name
             elif isinstance(zip_file, dict) and zip_file.get("name"):
                 zip_path = zip_file["name"]
+            if zip_path:
+                extracted, logs = extract_zip_list(zip_path, final_zip_password)
+                final_paths = extracted
 
-        if use_default_zip_pass and (not zip_password or zip_password.strip() == ""):
-            final_zip_password = default_zip_password
-        else:
-            final_zip_password = zip_password
-
-        adv = {}
+    # call core generator
+    adv = {}
+    device = None
+    if device_choice and device_choice != "auto":
+        device = device_choice  # 'cpu' or 'cuda'
 
-        for logs, transcripts, word_path, percent in transcribe_multiple(
-            audio_input,
+    try:
+        for logs_text, transcripts_text, word_path, percent in transcribe_multiple(
+            final_paths,
             model_name,
             adv,
             merge_checkbox=merge,
-            zip_file=zip_path,
-            zip_password=final_zip_password,
             enable_memory=enable_memory,
+            device=device,
         ):
-            yield logs, transcripts, word_path, percent
-
+            yield logs_text, transcripts_text, word_path, percent
     except Exception:
         tb = traceback.format_exc()
-        logs = f"EXCEPTION in run_transcription_wrapper:\n{tb}"
-        transcripts = "ERROR: transcription did not start or failed unexpectedly."
-        yield logs, transcripts, None, 100
+        logs_text = f"EXCEPTION in run_transcription_ui:\n{tb}"
+        transcripts_text = "ERROR: transcription did not start or failed unexpectedly."
+        yield logs_text, transcripts_text, None, 100
 
 
+# Build UI (Tabs)
 print("DEBUG: building Gradio Blocks", flush=True)
-
-with gr.Blocks(title="Whisper Transcriber") as demo:
+with gr.Blocks(title="Whisper Transcriber — Multi-tab") as demo:
     gr.Markdown(
-        "## Whisper Transcriber\n"
-        "Upload audio files or a ZIP on the left and click **Transcribe**.\n"
-        "Transcript, progress, download, and logs appear on the right."
+        "<h2>Whisper Transcriber</h2>"
+        "<p>Upload audio files or a ZIP, extract and choose files, then transcribe.</p>",
     )
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Input")
-
-            file_input = gr.File(
-                label="Audio files",
-                file_count="multiple",
-                type="filepath",
-                height=60,
-            )
-
-            zip_input = gr.File(
-                label="ZIP with audio (optional)",
-                file_count="single",
-                type="filepath",
-                height=60,
-            )
-
-            use_default_zip_pass = gr.Checkbox(
-                label="Use default ZIP password",
-                value=False,
-            )
-
-            default_zip_password = gr.Textbox(
-                label="Default ZIP password",
-                value="dietcoke1",
-                interactive=True,
-            )
-
-            zip_password = gr.Textbox(
-                label="ZIP password (override)",
-                placeholder="If empty, default password will be used",
-            )
-
-            model_select = gr.Dropdown(
-                choices=["small", "medium", "large", "base"],
-                value="small",
-                label="Whisper model",
-            )
-
-            merge_checkbox = gr.Checkbox(
-                label="Merge all transcripts into one .docx",
-                value=True,
-            )
-
-            memory_checkbox = gr.Checkbox(
-                label="Enable correction memory (use during transcription)",
-                value=False,
-            )
-
-            submit = gr.Button("Transcribe", variant="primary")
-
-            gr.Markdown("### Memory management")
-            mem_upload = gr.File(label="Import memory file (JSON or text)", file_count="single", type="file")
-            mem_import_btn = gr.Button("Import Memory File")
-            mem_manual_entry = gr.Textbox(label="Add word/phrase to memory (manual)", placeholder="Type a word or phrase")
-            mem_add_btn = gr.Button("Add to Memory")
-            mem_clear_btn = gr.Button("Clear Memory")
-            mem_view_btn = gr.Button("View Memory")
-            mem_status = gr.Textbox(label="Memory status", interactive=False, lines=6)
-
-        with gr.Column(scale=1):
-            gr.Markdown("### Output")
-
-            transcripts_out = gr.Textbox(
-                label="Transcript",
-                lines=18,
-                interactive=False,
-            )
-
-            progress_num = gr.Slider(
-                minimum=0,
-                maximum=100,
-                value=0,
-                step=1,
-                label="Progress (%)",
-                interactive=False,
+    with gr.Tabs():
+        # ---------------- Transcribe Tab ----------------
+        with gr.TabItem("Transcribe"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("### Inputs")
+
+                    file_input = gr.File(label="Audio files (optional)", file_count="multiple", type="filepath", height=80)
+                    zip_input = gr.File(label="ZIP with audio (optional)", file_count="single", type="filepath", height=80)
+
+                    with gr.Row():
+                        zip_password = gr.Textbox(label="ZIP password (override)", placeholder="Optional")
+                        use_default_zip_pass = gr.Checkbox(label="Use default ZIP password", value=False)
+                        default_zip_password = gr.Textbox(label="Default ZIP password", value="", interactive=True)
+
+                    model_select = gr.Dropdown(choices=["small", "medium", "large", "base"], value="small", label="Whisper model")
+                    device_choice = gr.Dropdown(choices=["auto", "cpu", "cuda"], value="auto", label="Device (auto tries default)")
+
+                    merge_checkbox = gr.Checkbox(label="Merge all transcripts into one .docx", value=True)
+                    memory_checkbox = gr.Checkbox(label="Enable correction memory", value=False)
+
+                    gr.Markdown("### ZIP extraction & file selection")
+                    extract_btn = gr.Button("Extract ZIP & List Files")
+                    extracted_files_check = gr.CheckboxGroup(choices=[], label="Select extracted files to transcribe (optional)", interactive=True)
+                    extract_logs = gr.Textbox(label="Extraction logs", interactive=False, lines=6)
+
+                    # action buttons
+                    transcribe_btn = gr.Button("Transcribe Selected / Uploaded")
+                with gr.Column(scale=1):
+                    gr.Markdown("### Output")
+                    transcripts_out = gr.Textbox(label="Transcript", lines=20, interactive=False)
+                    progress_num = gr.Slider(minimum=0, maximum=100, value=0, step=1, label="Progress (%)", interactive=False)
+                    download_file = gr.File(label="Merged .docx (when available)")
+                    logs = gr.Textbox(label="Logs", lines=12, interactive=False)
+
+            # Wire extract button
+            def _extract_click(zip_file, zip_password, use_default_zip_pass, default_zip_password):
+                choices, logstxt = extract_zip_for_ui(zip_file, zip_password, use_default_zip_pass, default_zip_password)
+                # choices are paths; show them in CheckboxGroup
+                return choices, logstxt
+
+            extract_btn.click(fn=_extract_click, inputs=[zip_input, zip_password, use_default_zip_pass, default_zip_password], outputs=[extracted_files_check, extract_logs])
+
+            # Wire transcribe button: need to pass selected extracted files (list), file_input, model, merge, zip file (for fallback), etc.
+            transcribe_btn.click(
+                fn=run_transcription_ui,
+                inputs=[
+                    extracted_files_check,  # zip_selected_files
+                    file_input,             # file_input
+                    model_select,
+                    merge_checkbox,
+                    # pass in zip file so fallback is possible
+                    extracted_files_check,  # placeholder to keep ordering (not used) - we will also pass zip_input below
+                    zip_input,
+                    zip_password,
+                    use_default_zip_pass,
+                    default_zip_password,
+                    memory_checkbox,
+                    device_choice,
+                ],
+                outputs=[logs, transcripts_out, download_file, progress_num],
             )
 
-            download_file = gr.File(
-                label="Merged .docx (when available)"
-            )
-
-            logs = gr.Textbox(
-                label="Logs",
-                lines=10,
-                interactive=False,
-            )
-
-    submit.click(
-        fn=run_transcription_wrapper,
-        inputs=[
-            file_input,
-            model_select,
-            merge_checkbox,
-            zip_input,
-            zip_password,
-            use_default_zip_pass,
-            default_zip_password,
-            memory_checkbox,
-            gr.State({}),
-        ],
-        outputs=[logs, transcripts_out, download_file, progress_num],
-    )
-
-    def _import_memory(uploaded):
-        return import_memory_file(uploaded)
-
-    mem_import_btn.click(fn=_import_memory, inputs=[mem_upload], outputs=[mem_status])
-
-    mem_add_btn.click(fn=add_memory_entry, inputs=[mem_manual_entry], outputs=[mem_status])
-
-    mem_clear_btn.click(fn=lambda: clear_memory(), inputs=[], outputs=[mem_status])
-
-    mem_view_btn.click(fn=lambda: view_memory(), inputs=[], outputs=[mem_status])
-
+        # ---------------- Memory Tab ----------------
+        with gr.TabItem("Memory"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("### Memory Tools")
+                    mem_upload = gr.File(label="Import memory file (JSON or text)", file_count="single", type="filepath")
+                    mem_import_btn = gr.Button("Import Memory File")
+                    mem_manual_entry = gr.Textbox(label="Add word/phrase to memory (manual)", placeholder="Type a word or phrase")
+                    mem_add_btn = gr.Button("Add to Memory")
+                    mem_clear_btn = gr.Button("Clear Memory")
+                    mem_view_btn = gr.Button("View Memory")
+                    mem_status = gr.Textbox(label="Memory status", interactive=False, lines=12)
+
+            # memory bindings
+            def _import_mem(uploaded):
+                return import_memory_file(uploaded)
+
+            mem_import_btn.click(fn=_import_mem, inputs=[mem_upload], outputs=[mem_status])
+            mem_add_btn.click(fn=add_memory_entry, inputs=[mem_manual_entry], outputs=[mem_status])
+            mem_clear_btn.click(fn=lambda: clear_memory(), inputs=[], outputs=[mem_status])
+            mem_view_btn.click(fn=lambda: view_memory(), inputs=[], outputs=[mem_status])
+
+        # ---------------- Settings Tab ----------------
+        with gr.TabItem("Settings"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Settings")
+                    gr.Markdown("- Use `Device` in Transcribe tab to force CPU/GPU. Default uses whisper's choice.")
+                    gr.Markdown("- `Default ZIP password` is empty by default for safety.")
+                    gr.Markdown("- If you want extracted-file preview before transcribing, click **Extract ZIP & List Files** first.")
+                with gr.Column():
+                    gr.Markdown("### Diagnostics")
+                    diag_btn = gr.Button("Show memory summary")
+                    diag_out = gr.Textbox(label="Diagnostics output", interactive=False, lines=12)
+
+                    diag_btn.click(fn=lambda: view_memory(), inputs=[], outputs=[diag_out])
+
+    # end tabs
+
+# ---------- Launch ----------
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
     print("DEBUG: launching Gradio on port", port, flush=True)