app.py

"""
World-class Speech Therapy Pipeline (end-to-end)
- Dataset: speechocean762 (used for fine-tuning ASR / optional model training)
- ASR: Whisper (recommended) OR Wav2Vec2 fine-tune (optional)
- Pronunciation checking: phoneme-level using g2p (g2p_en) + alignment
- Error detection: rule-based + optional classifier training placeholder
- Feedback: LLM (google/flan-t5-base by default) to explain issues & give remedies
- UI: Gradio for live demo

HOW TO USE:
1) Install required packages (see INSTALL block)
2) (Optional) Train ASR or Pronunciation classifier — heavy, needs GPU
3) Use the inference pipeline to analyze child speech and get detailed feedback

Run in Google Colab for easiest GPU access.
"""

# --------------------------
# INSTALL (run once)
# --------------------------
# In a notebook or terminal run:
# !pip install -q datasets soundfile librosa transformers accelerate \
#     jiwer gradio faster-whisper openai-whisper g2p-en phonemizer python-Levenshtein
# # ffmpeg install (Colab already has ffmpeg), else:
# # Ubuntu: sudo apt-get update && sudo apt-get install -y ffmpeg
# --------------------------

import nltk
nltk.download('averaged_perceptron_tagger_eng')
nltk.download('punkt')


import os
import json
import tempfile
import warnings
import subprocess
from difflib import SequenceMatcher
from typing import Tuple, List, Dict, Optional

import numpy as np
import soundfile as sf
import librosa
import gradio as gr

warnings.filterwarnings("ignore")

# Hugging Face & ASR imports (may be heavy)
try:
    from datasets import load_dataset, load_metric
    HF_DATASETS = True
except Exception:
    HF_DATASETS = False

try:
    import torch
    HAS_TORCH = True
except Exception:
    HAS_TORCH = False

# prefer faster-whisper if available for best Whisper inference
try:
    from faster_whisper import WhisperModel
    HAS_FASTER_WHISPER = True
except Exception:
    HAS_FASTER_WHISPER = False

try:
    import whisper as openai_whisper
    HAS_OPENAI_WHISPER = True
except Exception:
    HAS_OPENAI_WHISPER = False

# transformers for training/inference LLM and Wav2Vec2 if needed
try:
    import transformers
    from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
    TRANSFORMERS = True
except Exception:
    TRANSFORMERS = False

# phoneme tools
try:
    # g2p_en for English grapheme->phoneme
    from g2p_en import G2p
    G2P_AVAILABLE = True
except Exception:
    G2P_AVAILABLE = False

# Helpful util: Levenshtein distance if installed (fast)
try:
    import Levenshtein
    HAVE_LEV = True
except Exception:
    HAVE_LEV = False

# --------------------------
# Utilities: audio handling
# --------------------------
def ensure_ffmpeg():
    """Check ffmpeg present; if not, raise a readable error."""
    try:
        subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
    except Exception as e:
        raise RuntimeError("ffmpeg not found. Install ffmpeg (apt-get install ffmpeg) or run in Colab which has ffmpeg.")

def convert_to_wav(path_in: str, sr=16000) -> str:
    """
    Convert arbitrary audio (m4a/mp3) to mono 16k wav using ffmpeg.
    Returns path to wav file.
    """
    ensure_ffmpeg()
    base = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    cmd = ["ffmpeg", "-y", "-i", path_in, "-ar", str(sr), "-ac", "1", base]
    subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
    return base

def read_audio_any(path_or_np):
    """
    Accepts Gradio audio (filepath), or numpy array; returns (audio_np, sr)
    """
    if path_or_np is None:
        raise RuntimeError("No audio provided.")
    if isinstance(path_or_np, str):
        ext = os.path.splitext(path_or_np)[1].lower()
        if ext not in [".wav", ".flac", ".ogg"]:
            path_or_np = convert_to_wav(path_or_np, sr=16000)
        arr, sr = sf.read(path_or_np)
        if arr.ndim > 1:
            arr = np.mean(arr, axis=1)
        return arr.astype(np.float32), int(sr)
    if isinstance(path_or_np, np.ndarray):
        arr = path_or_np
        if arr.ndim > 1:
            arr = np.mean(arr, axis=1)
        return arr.astype(np.float32), 16000
    raise RuntimeError("Unsupported audio input type.")

def save_wav_from_array(arr: np.ndarray, sr=16000) -> str:
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(tmp.name, arr, sr)
    tmp.close()
    return tmp.name

# --------------------------
# ASR: Whisper inference wrapper (fast, offline)
# --------------------------
_whisper_fast_model = None
_openai_whisper_model = None
def load_whisper_model(size="small"):
    global _whisper_fast_model, _openai_whisper_model
    if HAS_FASTER_WHISPER:
        if _whisper_fast_model is None:
            device = "cuda" if (HAS_TORCH and torch.cuda.is_available()) else "cpu"
            # If CPU, use float32 to avoid float16 errors. If GPU, float16 is fine.
            compute_type = "float32" if device != "cpu" else "float32"
            _whisper_fast_model = WhisperModel(size, device=device, compute_type=compute_type)
        return ("faster-whisper", _whisper_fast_model)
    if HAS_OPENAI_WHISPER:
        if _openai_whisper_model is None:
            _openai_whisper_model = openai_whisper.load_model(size)
        return ("openai-whisper", _openai_whisper_model)
    if TRANSFORMERS:
        # fallback to HF pipeline - use english model variant .en to avoid language detection issues
        pipe = pipeline("automatic-speech-recognition", model=f"openai/whisper-{size}.en")
        return ("hf-whisper", pipe)
    raise RuntimeError("No ASR backend available. Install faster-whisper, openai-whisper, or transformers.")

def asr_transcribe_file(path: str, model_size="small", force_lang="en") -> Tuple[str, float]:
    """
    Transcribe file path -> (text, duration)
    force_lang: sets language to 'en' where supported to avoid auto-detect switching to Hindi etc.
    """
    backend, model = load_whisper_model(model_size)
    if backend == "faster-whisper":
        segments, info = model.transcribe(path, beam_size=5, language=force_lang)
        text = " ".join([s.text.strip() for s in segments])
        duration = getattr(info, "duration", librosa.get_duration(filename=path))
        return text.strip(), duration
    if backend == "openai-whisper":
        r = model.transcribe(path, language=force_lang)
        text = r.get("text", "").strip()
        dur = librosa.get_duration(filename=path)
        return text, dur
    if backend == "hf-whisper":
        out = model(path)
        # pipeline returns dict or string
        text = out.get("text", "") if isinstance(out, dict) else str(out)
        return text.strip(), librosa.get_duration(filename=path)
    raise RuntimeError("ASR backend failed")

# --------------------------
# Phoneme utilities (G2P)
# --------------------------
if G2P_AVAILABLE:
    g2p = G2p()
else:
    g2p = None

def text_to_phonemes(text: str) -> List[str]:
    """
    Convert English text to phoneme list using g2p_en (if available).
    Falls back to a simple char-level approximation if g2p not installed.
    """
    if not text:
        return []
    if g2p is None:
        # fallback: split characters (not ideal)
        return list(text.lower().replace(" ", ""))
    phones = g2p(text)
    # g2p returns list with spaces and punctuation; filter
    phones = [p for p in phones if p.strip() and p != " "]
    # sometimes g2p returns tokens and stress markers; keep them as simple tokens
    return phones

def phoneme_align(ref_phones: List[str], hyp_phones: List[str]) -> Dict:
    """
    Align two phoneme sequences using SequenceMatcher and return mismatches.
    Returns:
      - overall_similarity (0-100)
      - mispronounced list (phonemes from ref that mismatch)
      - alignment ops and details
    """
    matcher = SequenceMatcher(None, ref_phones, hyp_phones)
    ops = []
    mis = []
    sims = []
    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
        if tag == "equal":
            for k in range(i2 - i1):
                sims.append(1.0)
                ops.append(("equal", ref_phones[i1+k], hyp_phones[j1+k]))
        elif tag == "replace":
            # pair replacements
            for idx_ref in range(i1, i2):
                refp = ref_phones[idx_ref]
                recp = hyp_phones[j1] if j1 < j2 else ""
                sims.append(SequenceMatcher(None, refp, recp).ratio())
                ops.append(("replace", refp, recp))
                mis.append(refp)
                j1 += 1
        elif tag == "delete":
            for idx_ref in range(i1, i2):
                refp = ref_phones[idx_ref]
                ops.append(("delete", refp, ""))
                sims.append(0.0)
                mis.append(refp)
        elif tag == "insert":
            for idx_rec in range(j1, j2):
                recp = hyp_phones[idx_rec]
                ops.append(("insert", "", recp))
                # insertions do not penalize ref-sim directly
    overall = round(np.mean(sims) * 100.0, 2) if sims else 0.0
    return {"overall_similarity": overall, "mispronounced": mis, "ops": ops}

# --------------------------
# Error type heuristic classifier (rule-based) + placeholder ML training
# --------------------------
def error_type_rules(ref_word: str, hyp_word: str, ref_ph: List[str], hyp_ph: List[str]) -> str:
    """
    Heuristic rules for common error types:
      - substitution (one phoneme replaced)
      - omission (phoneme deleted)
      - insertion (extra phoneme)
      - distortion (reduced similarity)
    This is a simple rule-based classifier; you can replace with an ML model if you have labeled data.
    """
    alignment = phoneme_align(ref_ph, hyp_ph)
    mis = alignment["mispronounced"]
    if not mis:
        return "correct"
    # heuristics
    ops = alignment["ops"]
    for op in ops:
        if op[0] == "delete":
            return "omission"
        if op[0] == "insert":
            return "insertion"
        if op[0] == "replace":
            # if similarity low => distortion/substitution
            sim = SequenceMatcher(None, op[1], op[2]).ratio()
            if sim < 0.6:
                return "substitution"
            else:
                return "distortion"
    return "unknown"

# --------------------------
# LLM feedback generator (Flan-T5 default)
# --------------------------
_llm_pipe = None
def ensure_llm(model_name="google/flan-t5-base"):
    global _llm_pipe
    if _llm_pipe is None:
        if not TRANSFORMERS:
            raise RuntimeError("Transformers not installed; cannot load LLM.")
        # text2text for flan-t5; device mapping
        device = 0 if (HAS_TORCH and torch.cuda.is_available()) else -1
        _llm_pipe = pipeline("text2text-generation", model=model_name, device=device)
    return _llm_pipe

def generate_explanation_feedback(expected_text: str, recognized_text: str, per_word_report: List[Dict],
                                  model_name="google/flan-t5-base", max_length=250) -> str:
    """
    Build a prompt describing each problem word, error type, phonetic mismatch, and ask the LLM to give:
      - friendly explanation
      - 1-2 concrete exercises
      - 1 quick practice phrase
    """
    pipe = ensure_llm(model_name)
    # build concise prompt
    prompt_lines = [
        "You are a kind, patient speech therapist who explains clearly to parents and children.",
        "Given the expected sentence and the child's spoken sentence, explain what went wrong for each problematic word and give simple exercises."
    ]
    prompt_lines.append(f"Expected: {expected_text}")
    prompt_lines.append(f"Recognized: {recognized_text}")
    if not per_word_report:
        prompt_lines.append("No problems detected. Give a quick encouraging message and one maintenance exercise.")
    else:
        prompt_lines.append("Problems (word -> detected issue -> error type -> suggestion seed):")
        for w in per_word_report:
            # w contains: ref_word, hyp_word, ref_ph, hyp_ph, error_type, similarity
            prompt_lines.append(f"- {w['ref_word']} -> said as '{w['hyp_word']}' -> {w['error_type']} (sim {w['similarity']}%)")
    prompt_lines.append("\nNow write a short, child-friendly explanation for each problem word (2-3 sentences), and give one simple practice tip and one practice phrase to repeat 5 times.")
    prompt = "\n".join(prompt_lines)

    out = pipe(prompt, max_length=max_length, do_sample=False)
    txt = out[0].get("generated_text", out[0].get("text", "")).strip()
    return txt

# --------------------------
# High-level analyzer (word-level + phoneme-level)
# --------------------------
def analyze_pronunciation(expected_sentence: str, audio_input, asr_model_size="small", llm_model="google/flan-t5-base"):
    """
    Main function you can call from UI:
      - transcribe audio (Whisper default)
      - split into words and map recognized words to expected words (basic alignment)
      - for each pair, compute phonemes and mismatch
      - detect error type, build per-word report
      - call LLM to get explanation + exercises
    Returns (Markdown_text, JSON report)
    """
    # Step 0: basic checks
    if not expected_sentence or not expected_sentence.strip():
        return "Please type the expected sentence.", ""

    # Step 1: parse audio
    try:
        audio_np, sr = read_audio_any(audio_input)
    except Exception as e:
        return f"Error reading audio: {e}", ""

    # ensure 16k
    if sr != 16000:
        audio_np = librosa.resample(audio_np.astype(float), orig_sr=sr, target_sr=16000)
        sr = 16000
    # save temp
    wav_path = save_wav_from_array(audio_np, sr)

    # Step 2: ASR transcription (force english)
    try:
        transcript, duration = asr_transcribe_file(wav_path, model_size=asr_model_size, force_lang="en")
    except Exception as e:
        return f"ASR transcribe error: {e}", ""

    # Step 3: basic word alignment (align expected words to recognized words)
    expected_words = [w.strip() for w in expected_sentence.strip().split() if w.strip()]
    rec_words = [w.strip() for w in transcript.strip().split() if w.strip()]

    # Simple word alignment via SequenceMatcher (pair expected->recognized roughly)
    matcher = SequenceMatcher(None, expected_words, rec_words)
    word_pairs = []  # list of tuples (ref_word, hyp_word)
    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
        if tag == "equal":
            for k in range(i2 - i1):
                word_pairs.append((expected_words[i1+k], rec_words[j1+k]))
        elif tag == "replace":
            # pair sequentially (could be different lengths)
            len_pair = max(i2 - i1, j1 - j2 + (j2 - j1))
            # iterate min pairs
            r_range = list(range(i1, i2))
            h_range = list(range(j1, j2))
            n = max(len(r_range), len(h_range))
            for idx in range(n):
                ref_w = expected_words[i1 + idx] if idx < len(r_range) else ""
                hyp_w = rec_words[j1 + idx] if idx < len(h_range) else ""
                word_pairs.append((ref_w, hyp_w))
        elif tag == "delete":
            for idx in range(i1, i2):
                word_pairs.append((expected_words[idx], ""))
        elif tag == "insert":
            for idx in range(j1, j2):
                word_pairs.append(("", rec_words[idx]))

    # Step 4: phoneme-level analysis and error detection
    per_word_report = []
    for ref_w, hyp_w in word_pairs:
        # get phonemes
        ref_ph = text_to_phonemes(ref_w) if ref_w else []
        hyp_ph = text_to_phonemes(hyp_w) if hyp_w else []
        align = phoneme_align(ref_ph, hyp_ph)
        error_type = error_type_rules(ref_w, hyp_w, ref_ph, hyp_ph)
        per_word_report.append({
            "ref_word": ref_w or "(missing)",
            "hyp_word": hyp_w or "(missing)",
            "ref_ph": ref_ph,
            "hyp_ph": hyp_ph,
            "similarity": align["overall_similarity"],
            "mispronounced": align["mispronounced"],
            "error_type": error_type,
            "alignment_ops": align["ops"]
        })

    # Step 5: Build high-level stats & JSON
    overall_sim = np.mean([w["similarity"] for w in per_word_report]) if per_word_report else 0.0
    report = {
        "expected": expected_sentence,
        "recognized": transcript,
        "duration": duration,
        "overall_similarity": float(overall_sim),
        "per_word": per_word_report
    }

    # Step 6: LLM explanation & remedies
    # Use the LLM to produce clear child-friendly feedback
    try:
        explanation = generate_explanation_feedback(expected_sentence, transcript, per_word_report, model_name=llm_model)
    except Exception as e:
        # fallback to simple rule-based messages
        explanation = "Could not generate LLM feedback (model load failed). Basic suggestions:\n"
        for w in per_word_report:
            if w["error_type"] != "correct":
                explanation += f"For '{w['ref_word']}' said as '{w['hyp_word']}': {w['error_type']}. Try repeating the word slowly.\n"
        if overall_sim > 90:
            explanation = "Great job! Most words sound correct. Keep practicing regularly."

    # Step 7: produce Markdown output for UI
    md = f"### Transcript\n`{transcript or '(no speech recognized)'}`\n\n"
    md += f"**Duration:** {duration:.2f}s  \n"
    md += f"**Overall similarity:** **{report['overall_similarity']:.2f}%**  \n\n"
    md += "### Word-level results\n"
    for w in per_word_report:
        status = "✅ correct" if w["error_type"] == "correct" else f"❌ {w['error_type']}"
        md += f"- **{w['ref_word']}** -> said as `{w['hyp_word']}` : {status} (sim {w['similarity']:.1f}%)\n"
        if w["mispronounced"]:
            md += f"  - phonemes expected: {w['ref_ph']}\n  - phonemes spoken: {w['hyp_ph']}\n"
    md += "\n### Therapist-style feedback (LLM)\n"
    md += explanation + "\n"

    return md, json.dumps(report, indent=2)

# --------------------------
# Gradio UI
# --------------------------
title = "🏆 Advanced Speech Therapy Assistant — Phoneme-aware + LLM Feedback"
desc = "Upload/record child speech, type the expected sentence (what child should say). The system will transcribe, detect pronunciation flaws, and give therapist-like remedies."

with gr.Blocks() as demo:
    gr.Markdown(f"# {title}\n{desc}")
    with gr.Row():
        with gr.Column(scale=2):
            expected = gr.Textbox(label="Expected Sentence (type what child should say)", value="Twinkle twinkle little star")
            audio_in = gr.Audio(sources=["microphone", "upload"], type="filepath", label="🎤 Upload/Record audio")
            asr_size = gr.Dropdown(["tiny", "small", "medium", "large-v2"], value="small", label="ASR Model size (Whisper)")
            llm_choice = gr.Dropdown(["google/flan-t5-base", "google/flan-t5-small"], value="google/flan-t5-base", label="LLM model for feedback")
            analyze_btn = gr.Button("Analyze Pronunciation")
        with gr.Column(scale=1):
            out_md = gr.Markdown()
            out_json = gr.Textbox(label="Raw JSON report", visible=False)
    analyze_btn.click(fn=analyze_pronunciation, inputs=[expected, audio_in, asr_size, llm_choice], outputs=[out_md, out_json])

if __name__ == "__main__":
    demo.launch(debug=True, inline=True)
    demo.launch(share=True)