Update app.py

app.py

@@ -1,62 +1,36 @@
 #!/usr/bin/env python3
 """
-Medical transcription service (whisper) with NumPy compatibility fixes.
+Quick-start multi-file transcription -> merged DOCX using Whisper (small by default).
 
-Behavior:
-- If invoked with an audio path argument, transcribes that file (CLI mode).
-- If run with no args, starts a small Flask server with POST /transcribe for uploads.
-This avoids the container failing immediately when no CLI args are provided.
-"""
+- GET  /        -> simple upload UI
+- GET  /health  -> liveness (always 200)
+- GET  /ready   -> readiness (503 until model loaded)
+- POST /transcribe -> accept files field (multiple) and return merged docx
 
-import sys
+This app loads the Whisper model in a background thread so the server becomes responsive fast.
+"""
 import os
-import json
+import sys
+import tempfile
+import threading
+import time
 import logging
 from pathlib import Path
-import tempfile
+from typing import List, Dict, Any
+
+from flask import Flask, request, jsonify, send_file, render_template
 
-# Basic logging
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("med-asr")
-
-# Check numpy version early
-try:
-    import numpy as np
-    nv = tuple(map(int, np.__version__.split(".")[:2]))
-    if nv[0] >= 2:
-        logger.warning("NumPy version is %s. Some extensions compiled against NumPy 1.x may fail. "
-                       "If you see runtime errors, pin numpy<2 in requirements and rebuild the image.", np.__version__)
-except Exception:
-    # keep going; the Dockerfile installs numpy<2 so this should not happen.
-    pass
-
-# Import heavy deps and provide helpful message if missing
-try:
-    import torch
-except ModuleNotFoundError:
-    sys.stderr.write("Missing dependency: PyTorch is not installed. For CPU install, add to requirements.txt:\n"
-                     "  --extra-index-url https://download.pytorch.org/whl/cpu\n"
-                     "  torch==2.2.0+cpu\n")
-    sys.exit(1)
-
-try:
-    import whisper
-    import librosa
-    import soundfile as sf
-    from rapidfuzz import process, fuzz
-except Exception as e:
-    logger.exception("Failed to import dependencies: %s", e)
-    raise
-
-from flask import Flask, request, jsonify
-
-# Config (tune these)
-MODEL_NAME = os.environ.get("WHISPER_MODEL", "large-v2")
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-SAMPLE_RATE = 16000
-BEAM_SIZE = 5
-TEMPERATURE = 0.0
+logger = logging.getLogger("med-asr-quick")
+
+# Config
+MODEL_NAME = os.environ.get("WHISPER_MODEL", "small")
+PORT = int(os.environ.get("PORT", 5000))
+MAX_FILES = 20
+ALLOWED_EXT = {".wav", ".mp3", ".m4a", ".flac", ".aac", ".ogg"}
 MED_VOCAB_PATH = Path("medical_vocab.txt")
+
+# Postprocessing maps
 ABBREV_MAP = {
     "h/o": "history of",
     "c/o": "complains of",
@@ -67,6 +41,12 @@ ABBREV_MAP = {
     "bp": "blood pressure",
 }
 
+# Readiness globals
+READY = False
+MODEL = None
+MODEL_LOCK = threading.Lock()
+
+# Lazy-load med vocab
 def load_med_vocab(path=MED_VOCAB_PATH):
     if not path.exists():
         return []
@@ -75,22 +55,17 @@ def load_med_vocab(path=MED_VOCAB_PATH):
 
 MED_VOCAB = load_med_vocab()
 
-def whisper_transcribe(model, audio_path, beam_size=BEAM_SIZE, temperature=TEMPERATURE):
-    return model.transcribe(
-        audio_path,
-        language="en",
-        task="transcribe",
-        beam_size=beam_size,
-        temperature=temperature,
-        word_timestamps=False,
-    )
-
-def expand_abbrev(text):
+# Postprocessing helpers
+def expand_abbrev(text: str) -> str:
     for k, v in ABBREV_MAP.items():
         text = text.replace(k, v).replace(k.upper(), v)
     return text
 
-def medical_lexicon_correct(text, score_cutoff=70):
+def medical_lexicon_correct(text: str, score_cutoff: int = 70) -> str:
+    try:
+        from rapidfuzz import process, fuzz
+    except Exception:
+        return text
     if not MED_VOCAB:
         return text
     words = text.split()
@@ -100,7 +75,7 @@ def medical_lexicon_correct(text, score_cutoff=70):
         matched = False
         for n in (4, 3, 2):
             if i + n <= len(words):
-                phrase = " ".join(words[i:i+n])
+                phrase = " ".join(words[i : i + n])
                 res = process.extractOne(phrase, MED_VOCAB, scorer=fuzz.token_sort_ratio, score_cutoff=score_cutoff)
                 if res:
                     cand, score = res[0], res[1]
@@ -119,7 +94,7 @@ def medical_lexicon_correct(text, score_cutoff=70):
         i += 1
     return " ".join(out_words)
 
-def apply_postprocessing(text):
+def apply_postprocessing(text: str) -> str:
     text = text.strip()
     text = expand_abbrev(text)
     text = medical_lexicon_correct(text)
@@ -129,89 +104,161 @@ def apply_postprocessing(text):
             text = text + "."
     return text
 
-def redact_phi(text):
+def redact_phi(text: str) -> str:
     import re
     text = re.sub(r"\b(\d{3}-\d{2}-\d{4})\b", "[REDACTED_SSN]", text)
     text = re.sub(r"\b(\d{2}\/\d{2}\/\d{4})\b", "[REDACTED_DATE]", text)
     text = re.sub(r"\b(patient|pt)\s+\d+\b", "[REDACTED_PATIENT_ID]", text, flags=re.IGNORECASE)
     return text
 
-def transcribe_file(model, audio_path, redact=True):
-    logger.info("Transcribing %s on device=%s model=%s", audio_path, DEVICE, MODEL_NAME)
+# Transcription using whisper
+def whisper_transcribe(model, audio_path: str, beam_size: int = 5, temperature: float = 0.0) -> Dict[str, Any]:
+    return model.transcribe(audio_path, language="en", task="transcribe", beam_size=beam_size, temperature=temperature)
+
+def transcribe_single(model, audio_path: str, redact: bool = True) -> Dict[str, Any]:
     raw = whisper_transcribe(model, audio_path)
     text = raw.get("text", "").strip()
     proc = apply_postprocessing(text)
     if redact:
         proc = redact_phi(proc)
-    return {"raw_text": text, "postprocessed_text": proc, "segments": raw.get("segments", []), "language": raw.get("language", None)}
-
-# Load model once (when server starts or when CLI invoked)
-def load_model(name=MODEL_NAME):
-    logger.info("Loading model %s on %s (this may take a while)...", name, DEVICE)
-    m = whisper.load_model(name, device=DEVICE)
-    logger.info("Model loaded.")
-    return m
-
-# CLI mode: if audio argument provided, transcribe and exit
-def cli_mode(model, argv):
-    import argparse
-    parser = argparse.ArgumentParser(description="Medical transcription using Whisper")
-    parser.add_argument("audio", help="Path to audio file to transcribe (wav, mp3, m4a...)")
-    parser.add_argument("--no-redact", action="store_true", help="Disable automatic PHI redaction")
-    parser.add_argument("--output", "-o", default="transcript.json", help="Output JSON file")
-    args = parser.parse_args(argv)
-
-    audio_path = args.audio
-    if not Path(audio_path).exists():
-        logger.error("Audio file not found: %s", audio_path)
-        sys.exit(2)
-
-    out = transcribe_file(model, audio_path, redact=not args.no_redact)
-    with open(args.output, "w", encoding="utf-8") as f:
-        json.dump(out, f, indent=2)
-    logger.info("Saved transcript to %s", args.output)
-    print(out["postprocessed_text"])
-
-# Server mode: small Flask app for uploads
-app = Flask(__name__)
-MODEL = None
+    return {"raw_text": text, "postprocessed_text": proc, "segments": raw.get("segments", [])}
+
+def make_docx(trans_results: List[Dict[str, Any]], out_path: str):
+    from docx import Document
+    from docx.shared import Pt
+    doc = Document()
+    doc.styles["Normal"].font.name = "Arial"
+    doc.styles["Normal"].font.size = Pt(11)
+    doc.add_heading("Merged Transcripts", level=1)
+    for r in trans_results:
+        filename = r.get("filename", "unknown")
+        doc.add_heading(f"{filename}", level=2)
+        meta = r.get("meta", {})
+        if meta:
+            doc.add_paragraph(f"Duration: {meta.get('duration', 'unknown')}s")
+        doc.add_heading("Postprocessed Transcript", level=3)
+        doc.add_paragraph(r.get("postprocessed_text", ""))
+        doc.add_heading("Raw Transcript", level=3)
+        doc.add_paragraph(r.get("raw_text", ""))
+        segments = r.get("segments", [])
+        if segments:
+            doc.add_heading("Segments (timestamps)", level=4)
+            for seg in segments:
+                start = seg.get("start", 0)
+                end = seg.get("end", 0)
+                text = seg.get("text", "").strip()
+                doc.add_paragraph(f"[{start:.2f} - {end:.2f}] {text}")
+        doc.add_page_break()
+    doc.save(out_path)
+
+def validate_and_save_files(files_list) -> List[Path]:
+    saved_paths = []
+    if not files_list:
+        return []
+    if len(files_list) > MAX_FILES:
+        raise ValueError(f"Too many files (max {MAX_FILES}).")
+    for f in files_list:
+        filename = f.filename
+        if not filename:
+            continue
+        ext = Path(filename).suffix.lower()
+        if ext not in ALLOWED_EXT:
+            raise ValueError(f"Unsupported file extension: {ext}")
+        fd, tmp = tempfile.mkstemp(suffix=ext)
+        os.close(fd)
+        f.save(tmp)
+        saved_paths.append(Path(tmp))
+    return saved_paths
+
+# Background model loader
+def _load_model_background(name=MODEL_NAME):
+    global MODEL, READY
+    try:
+        import torch
+        import whisper
+    except Exception as e:
+        logger.exception("Failed to import heavy libs: %s", e)
+        READY = False
+        return
+    try:
+        logger.info("Loading Whisper model %s in background...", name)
+        m = whisper.load_model(name, device="cuda" if torch.cuda.is_available() else "cpu")
+        with MODEL_LOCK:
+            MODEL = m
+            READY = True
+        logger.info("Model loaded.")
+    except Exception:
+        logger.exception("Background model loading failed.")
+        READY = False
+
+def start_background_loader():
+    t = threading.Thread(target=_load_model_background, daemon=True)
+    t.start()
+
+# Flask app
+app = Flask(__name__, template_folder="templates", static_folder="static")
+
+@app.route("/")
+def index():
+    return render_template("index.html", max_files=MAX_FILES)
 
-@app.route("/health", methods=["GET"])
+@app.route("/health")
 def health():
-    return jsonify({"ok": True, "model": MODEL_NAME, "device": DEVICE})
+    return jsonify({"ok": True})
+
+@app.route("/ready")
+def ready():
+    return jsonify({"ready": READY, "model": MODEL_NAME}), (200 if READY else 503)
 
 @app.route("/transcribe", methods=["POST"])
 def transcribe_endpoint():
-    if "file" not in request.files:
-        return jsonify({"error": "missing file in form-data (name=file)"}), 400
-    f = request.files["file"]
-    if f.filename == "":
-        return jsonify({"error": "empty filename"}), 400
-    fd, tmp = tempfile.mkstemp(suffix=Path(f.filename).suffix)
-    os.close(fd)
-    f.save(tmp)
+    global READY, MODEL
+    if not READY or MODEL is None:
+        return jsonify({"error": "model not ready"}), 503
+    files = request.files.getlist("files")
+    if not files:
+        return jsonify({"error": "no files uploaded (use form field name 'files')"}), 400
+    try:
+        saved_paths = validate_and_save_files(files)
+    except ValueError as e:
+        return jsonify({"error": str(e)}), 400
+
+    trans_results = []
     try:
-        out = transcribe_file(MODEL, tmp, redact=True)
-        return jsonify(out)
+        for p in saved_paths:
+            logger.info("Transcribing %s", p)
+            try:
+                trans = transcribe_single(MODEL, str(p), redact=True)
+                # duration meta
+                try:
+                    import soundfile as sf
+                    info = sf.info(str(p))
+                    duration = round(info.duration, 2) if info.duration else None
+                except Exception:
+                    duration = None
+                trans_results.append({"filename": p.name, "meta": {"duration": duration}, **trans})
+            except Exception:
+                logger.exception("Error transcribing %s", p)
+                trans_results.append({"filename": p.name, "error": "transcription failed"})
     finally:
-        try:
-            os.remove(tmp)
-        except Exception:
-            pass
-
-def main():
-    global MODEL
-    # If arguments present -> CLI mode
-    if len(sys.argv) > 1:
-        MODEL = load_model()
-        cli_mode(MODEL, sys.argv[1:])
-        return
+        for p in saved_paths:
+            try:
+                os.remove(p)
+            except Exception:
+                pass
+
+    fd, out_tmp = tempfile.mkstemp(suffix=".docx")
+    os.close(fd)
+    try:
+        make_docx(trans_results, out_tmp)
+    except Exception:
+        logger.exception("Failed to create docx")
+        return jsonify({"error": "failed to generate docx"}), 500
 
-    # No args -> server mode
-    MODEL = load_model()
-    # start server
-    logger.info("Starting server on 0.0.0.0:5000")
-    app.run(host="0.0.0.0", port=5000, debug=False)
+    return send_file(out_tmp, as_attachment=True, download_name="merged_transcripts.docx")
 
 if __name__ == "__main__":
-    main()
+    # If CLI args provided, not supported here; run server
+    start_background_loader()
+    logger.info("Starting server on 0.0.0.0:%d (model loads in background)", PORT)
+    app.run(host="0.0.0.0", port=PORT, debug=False)