Update app.py

app.py

@@ -1,6 +1,6 @@
 # app.py
-# Whisper Transcriber — Full corrected app.py (multi-tab, Audio Transcribe focused)
-# Requirements: gradio, whisper, pydub, pyzipper, python-docx, ffmpeg installed.
+# Improved Whisper Transcriber (per-file selection after unzip, model availability check, SRT export)
+# Requirements: gradio, whisper, pydub, pyzipper, python-docx, ffmpeg
 
 import os
 import sys
@@ -14,10 +14,10 @@ import re
 from difflib import get_close_matches
 from pathlib import Path
 
-# Force unbuffered output so container logs show prints immediately
+# Force unbuffered prints for logs
 os.environ["PYTHONUNBUFFERED"] = "1"
 
-print("DEBUG: app.py bootstrap starting", flush=True)
+print("DEBUG: improved app.py bootstrap starting", flush=True)
 
 # Third-party imports
 try:
@@ -43,10 +43,11 @@ FFMPEG_CANDIDATES = [
     ("mulaw", 8000, 1),
 ]
 MODEL_CACHE = {}
+EXTRACT_MAP = {}  # maps friendly basename -> full path (populated after unzip)
 FINETUNE_WORKDIR = os.path.join(tempfile.gettempdir(), "finetune_workdir")
 os.makedirs(FINETUNE_WORKDIR, exist_ok=True)
 
-# ---------- Helpers: Memory & Postprocessing ----------
+# ---------- Helpers: memory & postprocessing ----------
 def load_memory():
     try:
         if os.path.exists(MEMORY_FILE):
@@ -67,6 +68,7 @@ def load_memory():
         pass
     return mem
 
+
 def save_memory(mem):
     with MEMORY_LOCK:
         try:
@@ -75,6 +77,7 @@ def save_memory(mem):
         except Exception:
             traceback.print_exc()
 
+
 memory = load_memory()
 
 MEDICAL_ABBREVIATIONS = {
@@ -90,13 +93,13 @@ MEDICAL_ABBREVIATIONS = {
     "adm": "admit",
     "disch": "discharge",
 }
-
 DRUG_NORMALIZATION = {
     "metformin": "Metformin",
     "aspirin": "Aspirin",
     "amoxicillin": "Amoxicillin",
 }
 
+
 def expand_abbreviations(text):
     tokens = re.split(r"(\s+)", text)
     out = []
@@ -112,11 +115,13 @@ def expand_abbreviations(text):
             out.append(t)
     return "".join(out)
 
+
 def normalize_drugs(text):
     for k, v in DRUG_NORMALIZATION.items():
         text = re.sub(rf"\b{k}\b", v, text, flags=re.IGNORECASE)
     return text
 
+
 def punctuation_and_capitalization(text):
     text = text.strip()
     if not text:
@@ -132,6 +137,7 @@ def punctuation_and_capitalization(text):
             out.append(p)
     return "".join(out)
 
+
 def postprocess_transcript(text):
     if not text:
         return text
@@ -141,11 +147,13 @@ def postprocess_transcript(text):
     t = punctuation_and_capitalization(t)
     return t
 
+
 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()]
     return [w for w in words if w.strip()], sentences
 
+
 def update_memory_with_transcript(transcript):
     global memory
     words, sentences = extract_words_and_phrases(transcript)
@@ -161,6 +169,7 @@ def update_memory_with_transcript(transcript):
         if changed:
             save_memory(memory)
 
+
 def memory_correct_text(text, min_ratio=0.85):
     if not text or (not memory.get("words") and not memory.get("phrases")):
         return text
@@ -194,7 +203,8 @@ def memory_correct_text(text, min_ratio=0.85):
             corrected = re.sub(re.escape(phrase), phrase, corrected, flags=re.IGNORECASE)
     return corrected
 
-# ---------- File utilities ----------
+
+# ---------- File utils ----------
 def save_as_word(text, filename=None):
     if filename is None:
         filename = os.path.join(tempfile.gettempdir(), "merged_transcripts.docx")
@@ -203,6 +213,7 @@ def save_as_word(text, filename=None):
     doc.save(filename)
     return filename
 
+
 # ---------- Conversion helpers ----------
 def _ffmpeg_convert(input_path, out_path, fmt, sr, ch):
     try:
@@ -230,6 +241,7 @@ def _ffmpeg_convert(input_path, out_path, fmt, sr, ch):
             pass
         return False, str(e)
 
+
 def convert_to_wav_if_needed(input_path):
     input_path = str(input_path)
     lower = input_path.lower()
@@ -312,7 +324,38 @@ def convert_to_wav_if_needed(input_path):
 
     raise Exception(f"Could not convert file to WAV. Diagnostics saved to: {diag_log}")
 
-# ---------- Whisper model loader ----------
+
+# ---------- Whisper model loader & availability ----------
+def whisper_available_models():
+    """Return set of model names if whisper provides helper; otherwise conservative fallback."""
+    try:
+        # many whisper forks expose available_models()
+        models = whisper.available_models()
+        if isinstance(models, (list, tuple, set)):
+            return set(models)
+    except Exception:
+        pass
+    # fallback: offer the common set but note we can't verify at startup
+    return set(["tiny", "base", "small", "medium", "large", "large-v3"])
+
+
+AVAILABLE_MODEL_SET = whisper_available_models()
+
+
+def safe_model_choices(prefer_default="small"):
+    # hide entries not in AVAILABLE_MODEL_SET
+    base_choices = ["small", "medium", "large", "large-v3", "base", "tiny"]
+    choices = [m for m in base_choices if m in AVAILABLE_MODEL_SET]
+    if not choices:
+        choices = base_choices  # if we couldn't detect, still present choices
+    # ensure prefer_default exists
+    if prefer_default in choices:
+        default = prefer_default
+    else:
+        default = choices[0]
+    return choices, default
+
+
 def get_whisper_model(name, device=None):
     if name not in MODEL_CACHE:
         print(f"DEBUG: loading whisper model '{name}'", flush=True)
@@ -322,12 +365,44 @@ def get_whisper_model(name, device=None):
             else:
                 MODEL_CACHE[name] = whisper.load_model(name)
         except TypeError:
-            # some whisper versions don't accept device arg
             MODEL_CACHE[name] = whisper.load_model(name)
     return MODEL_CACHE[name]
 
-# ---------- ZIP extraction helper ----------
-def extract_zip_list(zip_file, zip_password):
+
+# ---------- SRT export ----------
+def segments_to_srt(segments):
+    """
+    segments: iterable of dicts with 'start','end','text' or whisper segments
+    returns srt_text
+    """
+    def fmt_time(t):
+        # t in seconds
+        h = int(t // 3600)
+        m = int((t % 3600) // 60)
+        s = int(t % 60)
+        ms = int((t - int(t)) * 1000)
+        return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
+
+    lines = []
+    for i, seg in enumerate(segments, start=1):
+        start = seg.get("start", 0)
+        end = seg.get("end", 0)
+        text = seg.get("text", "").strip()
+        lines.append(str(i))
+        lines.append(f"{fmt_time(start)} --> {fmt_time(end)}")
+        lines.append(text)
+        lines.append("")  # blank line
+    return "\n".join(lines)
+
+
+# ---------- ZIP extraction + mapping for UI ----------
+def extract_zip_and_map(zip_path, zip_password=None):
+    """
+    Extracts supported audio files into temp dir and builds EXTRACT_MAP mapping friendly basename -> full path.
+    Returns list of friendly basenames and log string.
+    """
+    global EXTRACT_MAP
+    EXTRACT_MAP = {}
     temp_extract_dir = os.path.join(tempfile.gettempdir(), "extracted_audio")
     try:
         if os.path.exists(temp_extract_dir):
@@ -336,382 +411,247 @@ def extract_zip_list(zip_file, zip_password):
             except Exception:
                 pass
         os.makedirs(temp_extract_dir, exist_ok=True)
-        extracted = []
         logs = []
-        with pyzipper.ZipFile(zip_file, "r") as zf:
+        with pyzipper.ZipFile(zip_path, "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"]
+                    logs.append("Warning: failed to set zip password (continuing).")
+            count = {}
+            supported = [".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.")
+                if ext.lower() not in supported:
+                    continue
+                try:
+                    zf.extract(info, path=temp_extract_dir)
+                except RuntimeError as e:
+                    logs.append(f"Password required or incorrect for {info.filename}: {e}")
+                    continue
+                except Exception as e:
+                    logs.append(f"Error extracting {info.filename}: {e}")
+                    continue
+                fullp = os.path.normpath(os.path.join(temp_extract_dir, info.filename))
+                if not os.path.exists(fullp):
+                    continue
+                # friendly basename (avoid collisions)
+                base = os.path.basename(info.filename)
+                # if collision, append suffix
+                key = base
+                if key in EXTRACT_MAP:
+                    # create unique by adding index
+                    idx = count.get(base, 1) + 1
+                    count[base] = idx
+                    name_only, extn = os.path.splitext(base)
+                    key = f"{name_only} ({idx}){extn}"
+                else:
+                    count[base] = 1
+                EXTRACT_MAP[key] = fullp
+                logs.append(f"Extracted: {info.filename} -> {key}")
+        if not EXTRACT_MAP:
+            logs.append("No supported audio files found in ZIP.")
             return [], "\n".join(logs)
-        return extracted, "\n".join(logs)
+        # return sorted friendly names
+        friendly = sorted(EXTRACT_MAP.keys())
+        return friendly, "\n".join(logs)
     except Exception as e:
         traceback.print_exc()
         return [], f"Extraction failed: {e}"
 
-# ---------- Simple single-file transcriber ----------
-def transcribe_single(audio_path, model_name="small", enable_memory=False, device_choice="auto"):
+
+# ---------- Single-file transcribe (with SRT option) ----------
+def transcribe_single_file(path, model_name="small", device_choice="auto", enable_memory=False, generate_srt=False):
     logs = []
-    transcript_text = ""
     try:
-        if not audio_path:
-            return None, "No audio provided.", "No file provided."
-        path = str(audio_path)
+        if not path:
+            return None, "", "No file provided."
+        # normalize path if it's a file-like dict
+        p = path.name if hasattr(path, "name") else str(path)
         device = None if device_choice == "auto" else device_choice
         model = get_whisper_model(model_name, device=device)
         logs.append(f"Loaded model: {model_name}")
-        wav = convert_to_wav_if_needed(path)
+        wav = convert_to_wav_if_needed(p)
         logs.append(f"Converted to WAV: {os.path.basename(wav)}")
+        # call whisper transcribe
         result = model.transcribe(wav)
         text = result.get("text", "").strip()
         if enable_memory:
             text = memory_correct_text(text)
         text = postprocess_transcript(text)
-        transcript_text = text
+        srt_path = None
+        if generate_srt and result.get("segments"):
+            srt_text = segments_to_srt(result["segments"])
+            # save srt in temp dir
+            srt_fp = os.path.join(tempfile.gettempdir(), f"{os.path.splitext(os.path.basename(p))[0]}.srt")
+            with open(srt_fp, "w", encoding="utf-8") as fh:
+                fh.write(srt_text)
+            srt_path = srt_fp
+            logs.append(f"SRT generated: {srt_path}")
         if enable_memory:
             try:
                 update_memory_with_transcript(text)
                 logs.append("Memory updated.")
             except Exception:
                 pass
-        # cleanup temporary wav if created
-        if wav and os.path.exists(wav) and wav != path:
+        # cleanup intermediate wav if created
+        if wav and os.path.exists(wav) and wav != p:
             try:
                 os.unlink(wav)
             except Exception:
                 pass
-        return path, transcript_text, "\n".join(logs)
+        return text, srt_path, "\n".join(logs)
     except Exception as e:
         tb = traceback.format_exc()
-        return None, "", f"Error: {e}\n{tb}"
+        return "", None, f"Transcription error: {e}\n{tb}"
 
-# ---------- Fine-tune helpers (include old-files support) ----------
-def _collect_old_files_into(dst_dir, old_dir_path):
-    msgs = []
-    copied = 0
-    try:
-        if not os.path.isdir(old_dir_path):
-            return 0, f"Old-files path is not a directory: {old_dir_path}"
-        for root, _, files in os.walk(old_dir_path):
-            for f in files:
-                if f.lower().endswith((".wav", ".mp3", ".flac", ".m4a", ".ogg")):
-                    src_audio = os.path.join(root, f)
-                    base = os.path.splitext(f)[0]
-                    possible_txt = os.path.join(root, base + ".txt")
-                    rel_subdir = os.path.relpath(root, old_dir_path)
-                    target_subdir = os.path.join(dst_dir, rel_subdir)
-                    os.makedirs(target_subdir, exist_ok=True)
-                    target_audio = os.path.join(target_subdir, f)
-                    shutil.copy2(src_audio, target_audio)
-                    if os.path.exists(possible_txt):
-                        shutil.copy2(possible_txt, os.path.join(target_subdir, base + ".txt"))
-                        msgs.append(f"Copied pair: {os.path.join(rel_subdir, f)} + .txt")
-                    else:
-                        msgs.append(f"Copied audio (no transcript found): {os.path.join(rel_subdir, f)}")
-                    copied += 1
-        return copied, "\n".join(msgs)
-    except Exception as e:
-        traceback.print_exc()
-        return copied, f"Error copying old files: {e}"
 
-def prepare_finetune_dataset(uploaded_zip_or_dir, include_old_files=False, old_files_dir=""):
-    dst = os.path.join(FINETUNE_WORKDIR, "data")
-    try:
-        if os.path.exists(dst):
-            shutil.rmtree(dst)
-        os.makedirs(dst, exist_ok=True)
-    except Exception as e:
-        return f"Failed to prepare workdir: {e}", ""
-    path = None
-    try:
-        if uploaded_zip_or_dir:
-            if isinstance(uploaded_zip_or_dir, (str, os.PathLike)):
-                path = str(uploaded_zip_or_dir)
-            elif hasattr(uploaded_zip_or_dir, "name"):
-                path = uploaded_zip_or_dir.name
-            elif isinstance(uploaded_zip_or_dir, dict) and uploaded_zip_or_dir.get("name"):
-                path = uploaded_zip_or_dir["name"]
-    except Exception as e:
-        return f"Unable to determine uploaded path: {e}", ""
-    # extract or copy uploaded dataset if provided
-    if path and os.path.isfile(path) and path.lower().endswith(".zip"):
-        try:
-            with pyzipper.ZipFile(path, "r") as zf:
-                zf.extractall(dst)
-        except Exception as e:
-            return f"Failed to extract ZIP: {e}", ""
-    elif path and os.path.isdir(path):
-        try:
-            for item in os.listdir(path):
-                s = os.path.join(path, item)
-                d = os.path.join(dst, item)
-                if os.path.isdir(s):
-                    shutil.copytree(s, d)
-                else:
-                    shutil.copy2(s, d)
-        except Exception as e:
-            return f"Failed to copy dataset dir: {e}", ""
-    # include old files if requested
-    old_msgs = ""
-    if include_old_files and old_files_dir:
-        old_path = None
-        if isinstance(old_files_dir, (str, os.PathLike)):
-            old_path = str(old_files_dir)
-        elif hasattr(old_files_dir, "name"):
-            old_path = old_files_dir.name
-        elif isinstance(old_files_dir, dict) and old_files_dir.get("name"):
-            old_path = old_files_dir["name"]
-        if old_path:
-            copied, msg = _collect_old_files_into(dst, old_path)
-            old_msgs = f"\nOld-files: copied {copied} audio files.\nDetails:\n{msg}"
-    # find or build manifest
-    transcripts_candidates = [
-        os.path.join(dst, "transcripts.tsv"),
-        os.path.join(dst, "metadata.tsv"),
-        os.path.join(dst, "manifest.tsv"),
-        os.path.join(dst, "transcripts.txt"),
-        os.path.join(dst, "manifest.jsonl"),
-    ]
-    manifest_path = os.path.join(FINETUNE_WORKDIR, "manifest.tsv")
-    found = False
-    for tpath in transcripts_candidates:
-        if os.path.exists(tpath):
-            try:
-                shutil.copy2(tpath, manifest_path)
-                found = True
-                break
-            except Exception:
-                pass
-    missing_transcripts = 0
-    if not found:
-        audio_files = []
-        for root, _, files in os.walk(dst):
-            for f in files:
-                if f.lower().endswith((".wav", ".mp3", ".flac", ".m4a", ".ogg")):
-                    audio_files.append(os.path.join(root, f))
-        if not audio_files:
-            return f"No audio files found in dataset.{old_msgs}", ""
-        entries = []
-        for a in audio_files:
-            base = os.path.splitext(a)[0]
-            t_candidate = base + ".txt"
-            transcript = ""
-            if os.path.exists(t_candidate):
-                try:
-                    with open(t_candidate, "r", encoding="utf-8") as fh:
-                        transcript = fh.read().strip().replace("\n", " ")
-                except Exception:
-                    transcript = ""
+# ---------- Batch transcribe (maps friendly names to real paths) ----------
+def batch_transcribe(friendly_selected, uploaded_files, model_name, device_name, merge_flag, enable_mem, generate_srt):
+    logs = []
+    transcripts = []
+    srt_files = []
+    out_doc = None
+    paths = []
+    # selected from zip (friendly names)
+    if friendly_selected:
+        for key in friendly_selected:
+            p = EXTRACT_MAP.get(key)
+            if p:
+                paths.append(p)
             else:
-                missing_transcripts += 1
-            entries.append(f"{a}\t{transcript}")
+                logs.append(f"Warning: selected file not found in extract map: {key}")
+    # uploaded files
+    if uploaded_files:
+        if isinstance(uploaded_files, (list, tuple)):
+            for f in uploaded_files:
+                paths.append(str(f))
+        else:
+            paths.append(str(uploaded_files))
+    if not paths:
+        return "", "No files selected or uploaded.", None, None
+
+    total = len(paths)
+    for idx, p in enumerate(paths, start=1):
+        logs.append(f"[{idx}/{total}] Processing: {p}")
+        text, srt_path, lg = transcribe_single_file(p, model_name=model_name, device_choice=device_name, enable_memory=enable_mem, generate_srt=generate_srt)
+        logs.append(lg)
+        transcripts.append(f"FILE: {os.path.basename(p)}\n{text}\n")
+        if srt_path:
+            srt_files.append(srt_path)
+    combined = "\n\n".join(transcripts)
+    if merge_flag:
         try:
-            with open(manifest_path, "w", encoding="utf-8") as fh:
-                fh.write("\n".join(entries))
-            found = True
+            out_doc = save_as_word(combined)
+            logs.append(f"Merged transcript saved: {out_doc}")
         except Exception as e:
-            return f"Failed to write manifest: {e}{old_msgs}", ""
-    if not found:
-        return f"Failed to locate or build manifest.{old_msgs}", ""
-    status_msg = f"Dataset prepared. Manifest: {manifest_path}{old_msgs}"
-    if missing_transcripts > 0:
-        status_msg += f"\nWarning: {missing_transcripts} audio files have no matching .txt transcript (empty transcripts saved)."
-    return status_msg, manifest_path
-
-def start_finetune(manifest_path, base_model, epochs, batch_size, lr, output_dir):
-    outdir = output_dir or os.path.join(FINETUNE_WORKDIR, "output")
-    os.makedirs(outdir, exist_ok=True)
-    START_CMD = [
-        sys.executable,
-        "fine_tune.py",
-        "--manifest",
-        manifest_path,
-        "--base_model",
-        base_model,
-        "--epochs",
-        str(epochs),
-        "--batch_size",
-        str(batch_size),
-        "--lr",
-        str(lr),
-        "--output_dir",
-        outdir,
-    ]
-    try:
-        logfile = open(os.path.join(outdir, "finetune_stdout.log"), "a", encoding="utf-8")
-        proc = subprocess.Popen(START_CMD, stdout=logfile, stderr=logfile, cwd=os.getcwd())
-        return f"Fine-tune started (PID={proc.pid}). Logs: {logfile.name}"
-    except FileNotFoundError as e:
-        return f"Training script not found: {e}. Put 'fine_tune.py' in project root or change START_CMD."
-    except Exception as e:
-        return f"Failed to start fine-tune: {e}"
+            logs.append(f"Merge failed: {e}")
+    # if multiple SRTs, if desired we could zip them; here we just return first SRT if any
+    srt_return = srt_files[0] if srt_files else None
+    return combined, "\n".join(logs), out_doc, srt_return
 
-def tail_finetune_logs(logpath, lines=200):
-    try:
-        if not os.path.exists(logpath):
-            return "No logs yet."
-        with open(logpath, "r", encoding="utf-8", errors="ignore") as fh:
-            all_lines = fh.read().splitlines()
-            last = all_lines[-lines:]
-            return "\n".join(last)
-    except Exception as e:
-        return f"Failed to read logs: {e}"
 
-# ---------- UI CSS ----------
+# ---------- UI building ----------
+print("DEBUG: building Gradio UI", flush=True)
+
+available_choices, default_choice = safe_model_choices(prefer_default="small")
+
 CSS = """
-:root{
-  --accent:#4f46e5;
-  --muted:#6b7280;
-  --card:#ffffff;
-  --bg:#f7f8fb;
-}
+:root{ --accent:#4f46e5; --muted:#6b7280; --card:#ffffff; --bg:#f7f8fb; }
 body { background: var(--bg); font-family: Inter, system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", Arial; }
-.header { padding: 18px 24px; border-radius: 12px; background: linear-gradient(90deg, rgba(79,70,229,0.12), rgba(99,102,241,0.04)); margin-bottom: 18px; display:flex;align-items:center;gap:16px; }
-.app-icon { width:62px;height:62px;border-radius:12px;background:linear-gradient(135deg,var(--accent),#06b6d4);display:flex;align-items:center;justify-content:center;color:white;font-weight:700;font-size:24px; }
-.header-title h1 { margin:0;font-size:20px;}
-.header-sub { color:var(--muted); margin-top:4px;font-size:13px;}
-.card { background:var(--card); border-radius:12px; padding:14px; box-shadow: 0 6px 20px rgba(16,24,40,0.06); }
-.transcript-area { white-space:pre-wrap; font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, "Roboto Mono", monospace; background:#0f172a; color:#e6eef8; padding:12px; border-radius:10px; min-height:220px; }
+.header { padding: 14px; border-radius: 10px; background: linear-gradient(90deg, rgba(79,70,229,0.08), rgba(99,102,241,0.02)); margin-bottom: 12px; display:flex;align-items:center;gap:12px; }
+.app-icon { width:50px;height:50px;border-radius:10px;background:linear-gradient(135deg,var(--accent),#06b6d4);display:flex;align-items:center;justify-content:center;color:white;font-weight:700;font-size:20px; }
+.card { background:var(--card); border-radius:10px; padding:12px; box-shadow: 0 6px 20px rgba(16,24,40,0.04); }
+.transcript-area { white-space:pre-wrap; font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, "Roboto Mono", monospace; background:#0f172a; color:#e6eef8; padding:12px; border-radius:8px; min-height:200px; }
 .small-note { color:var(--muted); font-size:12px;}
 """
 
-# ---------- Build UI ----------
-print("DEBUG: building Gradio Blocks", flush=True)
-with gr.Blocks(title="Whisper Transcriber", css=CSS) as demo:
-    # Header
+with gr.Blocks(title="Whisper Transcriber (improved)", css=CSS) as demo:
     with gr.Row(elem_classes="header"):
         with gr.Column(scale=0):
             gr.HTML("<div class='app-icon'>WT</div>")
         with gr.Column():
-            gr.HTML("<h1 style='margin:0'>Whisper Transcriber</h1>")
-            gr.Markdown("<div class='header-sub'>Transcribe, batch, memory & fine-tune — multi-tab UI</div>")
+            gr.Markdown("<h3 style='margin:0'>Whisper Transcriber — improved</h3>")
+            gr.Markdown("<div class='small-note'>Per-file selection after unzip, SRT export, model availability checks.</div>")
 
     with gr.Tabs():
-        # Audio Transcribe Tab
+        # Single Audio Tab
         with gr.TabItem("Audio Transcribe"):
             with gr.Row():
                 with gr.Column(scale=1):
                     with gr.Group(elem_classes="card"):
-                        gr.Markdown("### Quick Single Audio Transcribe")
+                        gr.Markdown("### Single audio")
                         single_audio = gr.Audio(label="Upload or record audio", type="filepath")
                         with gr.Row():
-                            model_select = gr.Dropdown(choices=["small","medium","large","large-v3","base"], value="large-v3", label="Model")
-                            device_select = gr.Dropdown(choices=["auto","cpu","cuda"], value="auto", label="Device")
+                            model_select = gr.Dropdown(choices=available_choices, value=default_choice, label="Model")
+                            device_choice = gr.Dropdown(choices=["auto", "cpu", "cuda"], value="auto", label="Device")
                         with gr.Row():
-                            mem_toggle = gr.Checkbox(label="Enable correction memory", value=False)
-                            format_choice = gr.Dropdown(choices=["Plain","SOAP (medical)"], value="Plain", label="Format")
+                            mem_toggle = gr.Checkbox(label="Enable memory corrections", value=False)
+                            srt_toggle = gr.Checkbox(label="Generate SRT", value=False)
                         transcribe_btn = gr.Button("Transcribe", variant="primary")
-                        gr.Markdown("<div class='small-note'>Tip: choose large-v3 if your environment supports it.</div>")
                 with gr.Column(scale=1):
                     with gr.Group(elem_classes="card"):
-                        gr.Markdown("### Player & Transcript")
-                        audio_preview = gr.Audio(label="Player", interactive=False)
-                        transcript_out = gr.Textbox(label="Transcript", lines=14, interactive=False, elem_classes="transcript-area")
-                        transcript_logs = gr.Textbox(label="Logs", lines=6, interactive=False)
-
-            def _do_single_transcribe(audio_file, model_name, device_choice, enable_memory, fmt_choice):
-                player_path, transcript, logs = transcribe_single(audio_file, model_name=model_name, enable_memory=enable_memory, device_choice=device_choice)
-                if fmt_choice == "SOAP":
-                    sentences = re.split(r"(?<=[.?!])\s+", transcript)
-                    subj = sentences[0] if sentences else ""
-                    obj = sentences[1] if len(sentences) > 1 else ""
-                    soap = f"S: {subj}\nO: {obj}\nA: Assessment pending\nP: Plan: follow up"
-                    transcript = soap
-                return player_path, transcript, logs
-
-            transcribe_btn.click(fn=_do_single_transcribe, inputs=[single_audio, model_select, device_select, mem_toggle, format_choice], outputs=[audio_preview, transcript_out, transcript_logs])
-
-        # Batch Transcribe Tab
+                        gr.Markdown("### Output")
+                        audio_preview = gr.Audio(interactive=False)
+                        transcript_out = gr.Textbox(label="Transcript", lines=12, interactive=False, elem_classes="transcript-area")
+                        srt_download = gr.File(label="SRT (if generated / available)")
+                        single_logs = gr.Textbox(label="Logs", lines=8, interactive=False)
+
+            def _single_action(audio_file, model_name, device, mem_on, srt_on):
+                if not audio_file:
+                    return None, "", None, "No audio file provided."
+                path = audio_file if isinstance(audio_file, str) else (audio_file.name if hasattr(audio_file, "name") else str(audio_file))
+                text, srt_path, logs = transcribe_single_file(path, model_name=model_name, device_choice=device, enable_memory=mem_on, generate_srt=srt_on)
+                # set audio preview to original file
+                preview = audio_file
+                return preview, text, srt_path, logs
+
+            transcribe_btn.click(fn=_single_action, inputs=[single_audio, model_select, device_choice, mem_toggle, srt_toggle], outputs=[audio_preview, transcript_out, srt_download, single_logs])
+
+        # Batch Tab
         with gr.TabItem("Batch Transcribe"):
             with gr.Row():
                 with gr.Column(scale=1):
                     with gr.Group(elem_classes="card"):
-                        gr.Markdown("### Batch / ZIP workflow")
-                        batch_files = gr.File(label="Upload multiple audio files (optional)", file_count="multiple", type="filepath")
+                        gr.Markdown("### Batch (upload multiple or ZIP)")
+                        batch_files = gr.File(label="Upload audio files (optional)", file_count="multiple", type="filepath")
                         batch_zip = gr.File(label="Or upload ZIP with audio (optional)", file_count="single", type="filepath")
                         zip_password = gr.Textbox(label="ZIP password (optional)")
+                        batch_extract_btn = gr.Button("Extract ZIP & List files")
+                        batch_extract_logs = gr.Textbox(label="Extraction logs", lines=6, interactive=False)
+                        batch_select = gr.CheckboxGroup(choices=[], label="Select extracted files (friendly names)", interactive=True)
                         with gr.Row():
-                            batch_model = gr.Dropdown(choices=["small","medium","large","large-v3","base"], value="small", label="Model")
-                            batch_device = gr.Dropdown(choices=["auto","cpu","cuda"], value="auto", label="Device")
-                        batch_merge = gr.Checkbox(label="Merge all transcripts into one .docx", value=True)
+                            batch_model = gr.Dropdown(choices=available_choices, value=default_choice, label="Model")
+                            batch_device = gr.Dropdown(choices=["auto", "cpu", "cuda"], value="auto", label="Device")
+                        batch_merge = gr.Checkbox(label="Merge transcripts to DOCX", value=True)
                         batch_mem = gr.Checkbox(label="Enable memory corrections", value=False)
-                        batch_extract_btn = gr.Button("Extract ZIP & List Files")
-                        batch_extract_logs = gr.Textbox(label="Extraction logs", lines=6, interactive=False)
-                        batch_select = gr.CheckboxGroup(choices=[], label="Select extracted files to transcribe", interactive=True)
-                        batch_trans_btn = gr.Button("Start Batch Transcription", variant="primary")
+                        batch_srt = gr.Checkbox(label="Generate SRT(s) if available", value=False)
+                        batch_run_btn = gr.Button("Start Batch Transcription", variant="primary")
                 with gr.Column(scale=1):
                     with gr.Group(elem_classes="card"):
-                        gr.Markdown("### Output")
+                        gr.Markdown("### Batch Output")
                         batch_trans_out = gr.Textbox(label="Transcript (combined)", lines=16, interactive=False)
                         batch_logs = gr.Textbox(label="Logs", lines=10, interactive=False)
-                        batch_download = gr.File(label="Merged .docx (when available)")
+                        batch_doc_download = gr.File(label="Merged DOCX (if created)")
+                        batch_srt_download = gr.File(label="First SRT (if any)")
 
-            def _extract_zip_for_ui(zip_file, password):
+            def _do_extract(zip_file, password):
                 if not zip_file:
-                    return [], "No zip provided."
+                    return [], "No ZIP provided."
                 zip_path = zip_file.name if hasattr(zip_file, "name") else str(zip_file)
-                extracted, logs = extract_zip_list(zip_path, password)
-                short_logs = logs + "\n\nFiles:\n" + "\n".join([os.path.basename(p) for p in extracted])
-                return extracted, short_logs
-
-            batch_extract_btn.click(fn=_extract_zip_for_ui, inputs=[batch_zip, zip_password], outputs=[batch_select, batch_extract_logs])
-
-            def _batch_transcribe(selected_check, uploaded_files, model_name, device_name, merge_flag, enable_mem):
-                paths = []
-                if selected_check:
-                    paths.extend(selected_check)
-                if uploaded_files:
-                    if isinstance(uploaded_files, (list, tuple)):
-                        for x in uploaded_files:
-                            paths.append(str(x))
-                    else:
-                        paths.append(str(uploaded_files))
-                if not paths:
-                    return "", "No files selected or uploaded.", None
-                logs = []
-                transcripts = []
-                out_doc = None
-                for p in paths:
-                    try:
-                        _, txt, lg = transcribe_single(p, model_name=model_name, enable_memory=enable_mem, device_choice=device_name)
-                        logs.append(lg)
-                        transcripts.append(f"FILE: {os.path.basename(str(p))}\n{txt}\n")
-                    except Exception as e:
-                        logs.append(f"Failed {p}: {e}")
-                combined = "\n\n".join(transcripts)
-                if merge_flag:
-                    try:
-                        out_doc = save_as_word(combined)
-                        logs.append(f"Merged saved: {out_doc}")
-                    except Exception as e:
-                        logs.append(f"Merge failed: {e}")
-                return combined, "\n".join(logs), out_doc
+                friendly, logs = extract_zip_and_map(zip_path, password)
+                # Show friendly names and logs
+                return friendly, logs
+
+            batch_extract_btn.click(fn=_do_extract, inputs=[batch_zip, zip_password], outputs=[batch_select, batch_extract_logs])
+
+            def _do_batch(friendly_selected, uploaded_files, model_name, device, merge_flag, mem_flag, srt_flag):
+                combined, logs, out_doc, srt_path = batch_transcribe(friendly_selected, uploaded_files, model_name, device, merge_flag, mem_flag, srt_flag)
+                return combined, logs, out_doc, srt_path
 
-            batch_trans_btn.click(fn=_batch_transcribe, inputs=[batch_select, batch_files, batch_model, batch_device, batch_merge, batch_mem], outputs=[batch_trans_out, batch_logs, batch_download])
+            batch_run_btn.click(fn=_do_batch, inputs=[batch_select, batch_files, batch_model, batch_device, batch_merge, batch_mem, batch_srt], outputs=[batch_trans_out, batch_logs, batch_doc_download, batch_srt_download])
 
         # Memory Tab
         with gr.TabItem("Memory"):
@@ -719,9 +659,9 @@ with gr.Blocks(title="Whisper Transcriber", css=CSS) as demo:
                 with gr.Column(scale=1):
                     with gr.Group(elem_classes="card"):
                         gr.Markdown("### Correction Memory")
-                        mem_upload = gr.File(label="Import memory (JSON or text)", file_count="single", type="filepath")
+                        mem_upload = gr.File(label="Import memory file (JSON or text)", file_count="single", type="filepath")
                         mem_import_btn = gr.Button("Import Memory")
-                        mem_add_text = gr.Textbox(label="Add word / phrase", placeholder="Type and click Add")
+                        mem_text = gr.Textbox(label="Add word/phrase", placeholder="Type 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")
@@ -790,87 +730,39 @@ with gr.Blocks(title="Whisper Transcriber", css=CSS) as demo:
             def _view_mem():
                 w = memory.get("words", {})
                 p = memory.get("phrases", {})
-                out = []
-                out.append("WORDS (top 30):")
+                out_lines = []
+                out_lines.append("WORDS (top 30):")
                 for k, v in sorted(w.items(), key=lambda kv: -kv[1])[:30]:
-                    out.append(f"{k}: {v}")
-                out.append("")
-                out.append("PHRASES (top 20):")
+                    out_lines.append(f"{k}: {v}")
+                out_lines.append("")
+                out_lines.append("PHRASES (top 20):")
                 for k, v in sorted(p.items(), key=lambda kv: -kv[1])[:20]:
-                    out.append(f"{k}: {v}")
-                return "\n".join(out)
+                    out_lines.append(f"{k}: {v}")
+                return "\n".join(out_lines)
 
             mem_import_btn.click(fn=_import_mem, inputs=[mem_upload], outputs=[mem_status])
-            mem_add_btn.click(fn=_add_mem, inputs=[mem_add_text], outputs=[mem_status])
+            mem_add_btn.click(fn=_add_mem, inputs=[mem_text], outputs=[mem_status])
             mem_clear_btn.click(fn=_clear_mem, inputs=[], outputs=[mem_status])
             mem_view_btn.click(fn=_view_mem, inputs=[], outputs=[mem_status])
 
-        # Fine-tune Tab
-        with gr.TabItem("Fine-tune"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    with gr.Group(elem_classes="card"):
-                        gr.Markdown("### Prepare & Launch Fine-tune")
-                        ft_upload = gr.File(label="Upload dataset ZIP (optional)", file_count="single", type="filepath")
-                        ft_include_old = gr.Checkbox(label="Include old audio+transcript folder", value=False)
-                        ft_old = gr.File(label="Old files folder (optional)", file_count="single", type="filepath")
-                        ft_prepare_btn = gr.Button("Prepare dataset")
-                        ft_manifest_box = gr.Textbox(label="Prepare status / manifest", lines=4, interactive=False)
-                        ft_base_model = gr.Dropdown(choices=["small","base","medium","large","large-v3"], value="small", label="Base model")
-                        ft_epochs = gr.Slider(minimum=1, maximum=100, value=3, step=1, label="Epochs")
-                        ft_batch = gr.Number(label="Batch size", value=8)
-                        ft_lr = gr.Number(label="Learning rate", value=1e-5, precision=8)
-                        ft_output_dir = gr.Textbox(label="Output dir (optional)", value="", placeholder="Leave blank to use temp output")
-                        ft_start_btn = gr.Button("Start Fine-tune")
-                        ft_stop_btn = gr.Button("Stop Fine-tune")
-                        ft_start_status = gr.Textbox(label="Start/Stop status", interactive=False, lines=4)
-                        ft_tail_btn = gr.Button("Tail training logs")
-                        ft_logs = gr.Textbox(label="Training logs (tail)", interactive=False, lines=12)
-                with gr.Column(scale=1):
-                    with gr.Group(elem_classes="card"):
-                        gr.Markdown("### Notes")
-                        gr.Markdown("- Old-files folder should contain audio files and matching .txt transcripts with the same basename.")
-                        gr.Markdown("- The app prepares a manifest and calls your `fine_tune.py` training script (you must provide it).")
-
-            def _prepare_action(ft_upload_file, include_old, old_dir):
-                status, manifest = prepare_finetune_dataset(ft_upload_file, include_old_files=include_old, old_files_dir=old_dir)
-                return status
-
-            def _start_action(manifest_text, base_model, epochs, batch_size, lr, output_dir):
-                manifest_guess = os.path.join(FINETUNE_WORKDIR, "manifest.tsv")
-                if not os.path.exists(manifest_guess):
-                    return "Manifest not found. Prepare dataset first or manually provide manifest."
-                status = start_finetune(manifest_guess, base_model, int(epochs), int(batch_size), float(lr), output_dir)
-                return status
-
-            ft_prepare_btn.click(fn=_prepare_action, inputs=[ft_upload, ft_include_old, ft_old], outputs=[ft_manifest_box])
-            ft_start_btn.click(fn=_start_action, inputs=[ft_manifest_box, ft_base_model, ft_epochs, ft_batch, ft_lr, ft_output_dir], outputs=[ft_start_status])
-            ft_stop_btn.click(fn=lambda: "Stop not implemented in placeholder", inputs=[], outputs=[ft_start_status])
-            ft_tail_btn.click(fn=lambda: "Tail logs not implemented in placeholder", inputs=[], outputs=[ft_logs])
-
         # Settings Tab
         with gr.TabItem("Settings"):
             with gr.Row():
                 with gr.Column():
                     with gr.Group(elem_classes="card"):
                         gr.Markdown("### Runtime & tips")
-                        gr.Markdown("- Use large-v3 only if your whisper package supports it.")
+                        gr.Markdown("- Use `large-v3` only if your whisper package supports it.")
                         gr.Markdown("- Extraction writes to system temp `extracted_audio`. Re-extracting overwrites it.")
-                        gr.Markdown("- Provide your `fine_tune.py` for real fine-tuning.")
+                        gr.Markdown("- Provide `fine_tune.py` if you plan to use the Fine-tune workflow.")
                 with gr.Column():
                     with gr.Group(elem_classes="card"):
                         gr.Markdown("### Diagnostics")
                         diag_btn = gr.Button("Show memory summary")
                         diag_out = gr.Textbox(label="Diagnostics", lines=12, interactive=False)
-                        diag_btn.click(fn=lambda: _view_mem(), inputs=[], outputs=[diag_out])
+                        diag_btn.click(fn=lambda: (lambda: _view_mem())(), inputs=[], outputs=[diag_out])
 
-# ---------- Launch ----------
+# Launch
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    print("DEBUG: launching Gradio on port", port, flush=True)
-    try:
-        demo.queue().launch(server_name="0.0.0.0", server_port=port)
-    except Exception as e:
-        print("FATAL: demo.launch failed:", e, flush=True)
-        traceback.print_exc()
-        raise
+    print("DEBUG: launching improved Gradio on port", port, flush=True)
+    demo.queue().launch(server_name="0.0.0.0", server_port=port)