app.py

"""
AI Nexus · P1 demo — "Speech that stays yours"
A staged, visual walkthrough of the P1 assistive-speech pipeline:

    1 · Heard       — faster-whisper (ASR) turns degraded speech into a raw transcript
    2 · Understood  — a small local LLM corrects it into the intended message
    3 · Spoken      — Piper TTS says the clear message out loud

Plus a live "green meter" (energy + CO2 per run) via codecarbon.

This public Space runs in the cloud for convenience only.
The product itself runs the exact same pipeline fully on-device (Raspberry Pi 5).
"""

import difflib
import math
import os
import tempfile
import time
import wave

import gradio as gr
from faster_whisper import WhisperModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from piper import PiperVoice

# ----------------------------------------------------------------------------- 
# Configuration (override via Space "Variables" — no code changes needed)
# -----------------------------------------------------------------------------
WHISPER_MODEL = os.getenv("WHISPER_MODEL", "small")          # tiny/base/small
ASR_LANG = os.getenv("ASR_LANG", "en")

# Default LLM is Qwen 2.5 1.5B Instruct (Apache-2.0, no license gate).
# To use Llama 3.2 1B instead, set Space variables:
#   LLM_REPO = bartowski/Llama-3.2-1B-Instruct-GGUF
#   LLM_FILE = Llama-3.2-1B-Instruct-Q4_K_M.gguf
LLM_REPO = os.getenv("LLM_REPO", "Qwen/Qwen2.5-1.5B-Instruct-GGUF")
LLM_FILE = os.getenv("LLM_FILE", "qwen2.5-1.5b-instruct-q4_k_m.gguf")

# A UK voice to match the team; swap for any voice in rhasspy/piper-voices.
PIPER_ONNX = os.getenv(
    "PIPER_ONNX", "en/en_GB/alan/medium/en_GB-alan-medium.onnx"
)

CARBON_COUNTRY = os.getenv("CARBON_COUNTRY", "GBR")

# ----------------------------------------------------------------------------- 
# Load models once at startup
# -----------------------------------------------------------------------------
print(f"[startup] loading Whisper '{WHISPER_MODEL}' (CPU, int8)…")
ASR = WhisperModel(WHISPER_MODEL, device="cpu", compute_type="int8")

print(f"[startup] downloading LLM {LLM_REPO}/{LLM_FILE}…")
llm_path = hf_hub_download(repo_id=LLM_REPO, filename=LLM_FILE)
LLM = Llama(
    model_path=llm_path,
    n_ctx=2048,
    n_threads=max(2, os.cpu_count() or 2),
    verbose=False,
)

print(f"[startup] downloading Piper voice {PIPER_ONNX}…")
piper_model = hf_hub_download(repo_id="rhasspy/piper-voices", filename=PIPER_ONNX)
piper_cfg = hf_hub_download(
    repo_id="rhasspy/piper-voices", filename=PIPER_ONNX + ".json"
)  # downloaded so it sits next to the .onnx; Piper >=1.3 auto-detects it
TTS = PiperVoice.load(piper_model)

print("[startup] ready.")

SYSTEM_PROMPT = (
    "You are the language layer of an assistive communication aid. "
    "You receive the raw automatic transcript of speech from a person with a "
    "speech impairment. The transcript may contain recognition errors, missing "
    "words, or fragments. Rewrite it as the single most likely intended "
    "message, in the first person, preserving the speaker's meaning and tone. "
    "If the transcript is already clear, return it unchanged. "
    "Reply with the corrected sentence only — no commentary, no quotes."
)


# ----------------------------------------------------------------------------- 
# Pipeline stages
# -----------------------------------------------------------------------------
def stage_heard(audio_path: str, phrase_bank: str):
    """ASR with optional phrase-bank biasing (personalization preview)."""
    prompt = phrase_bank.strip()[:200] if phrase_bank else None
    segments, _info = ASR.transcribe(
        audio_path,
        language=ASR_LANG,
        beam_size=5,
        vad_filter=True,
        initial_prompt=prompt,
    )
    segs = list(segments)
    text = " ".join(s.text.strip() for s in segs).strip()
    if segs:
        conf = sum(math.exp(s.avg_logprob) for s in segs) / len(segs)
    else:
        conf = 0.0
    return text, conf


def stage_understood(raw: str) -> str:
    """Small-LLM correction into the intended message."""
    out = LLM.create_chat_completion(
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": raw},
        ],
        max_tokens=96,
        temperature=0.2,
    )
    text = out["choices"][0]["message"]["content"].strip().strip('"')
    return text or raw


def stage_spoken(text: str) -> str:
    """Piper TTS to a temp wav, returned as a filepath for gr.Audio."""
    path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    with wave.open(path, "wb") as f:
        TTS.synthesize_wav(text, f)
    return path


def diff_html(raw: str, fixed: str) -> str:
    """Word-level diff: what the corrector changed."""
    a, b = raw.split(), fixed.split()
    sm = difflib.SequenceMatcher(None, a, b)
    parts = []
    for op, i1, i2, j1, j2 in sm.get_opcodes():
        if op == "equal":
            parts.append(" ".join(a[i1:i2]))
        if op in ("replace", "delete") and i2 > i1:
            parts.append(f'<del>{" ".join(a[i1:i2])}</del>')
        if op in ("replace", "insert") and j2 > j1:
            parts.append(f'<ins>{" ".join(b[j1:j2])}</ins>')
    body = " ".join(p for p in parts if p)
    return f'<div class="diffbox"><span class="difflabel">What changed</span> {body}</div>'


# ----------------------------------------------------------------------------- 
# Orchestrator (generator → the stages light up one by one in the UI)
# -----------------------------------------------------------------------------
def run(audio_path, phrase_bank):
    if not audio_path:
        raise gr.Error("Record or upload a short clip first (a sentence or two).")

    tracker = None
    try:
        from codecarbon import OfflineEmissionsTracker

        tracker = OfflineEmissionsTracker(
            country_iso_code=CARBON_COUNTRY, log_level="error", save_to_file=False
        )
        tracker.start()
    except Exception:
        tracker = None  # meter is best-effort; never block the demo

    skip = gr.skip()

    # Stage 1 — Heard
    t0 = time.perf_counter()
    raw, conf = stage_heard(audio_path, phrase_bank)
    t_asr = time.perf_counter() - t0
    if not raw:
        if tracker:
            try:
                tracker.stop()
            except Exception:
                pass
        yield (
            "(no speech detected)",
            "Try again a little closer to the microphone.",
            skip, skip, skip,
            "—",
        )
        return
    conf_line = f"Recognition confidence ≈ {conf:.0%} · {t_asr:.1f}s"
    yield raw, conf_line, skip, skip, skip, "Stage 2 running…"

    # Stage 2 — Understood
    t0 = time.perf_counter()
    fixed = stage_understood(raw)
    t_llm = time.perf_counter() - t0
    yield raw, conf_line, fixed, diff_html(raw, fixed), skip, "Stage 3 running…"

    # Stage 3 — Spoken
    t0 = time.perf_counter()
    wav = stage_spoken(fixed)
    t_tts = time.perf_counter() - t0

    # Green meter
    meter = f"Heard {t_asr:.1f}s · Understood {t_llm:.1f}s · Spoken {t_tts:.1f}s"
    if tracker:
        try:
            kg = tracker.stop()
            data = tracker.final_emissions_data
            wh = (data.energy_consumed or 0) * 1000.0
            meter += f" &nbsp;|&nbsp; ⚡ {wh:.2f} Wh · 🌱 {kg * 1000:.2f} g CO₂e"
        except Exception:
            pass

    yield raw, conf_line, fixed, diff_html(raw, fixed), wav, meter


# ----------------------------------------------------------------------------- 
# UI
# -----------------------------------------------------------------------------
CSS = """
:root { --moss:#3F7D4E; --ink:#222A22; --paper:#F6F7F3; }
.gradio-container { max-width: 1080px !important; }
#hero h1 { font-family:'Fraunces', Georgia, serif; font-weight:600;
           letter-spacing:-0.01em; margin:0 0 4px 0; }
#hero .brand { text-transform:uppercase; letter-spacing:0.18em; font-size:12px;
               color:var(--moss); font-weight:700; }
#hero .sub { color:#5c665c; margin:0; }
.pipe { display:flex; gap:10px; align-items:center; margin:14px 0 4px 0;
        flex-wrap:wrap; }
.pipe .step { display:flex; gap:8px; align-items:center; background:#fff;
        border:1px solid #e2e6e0; border-radius:999px; padding:6px 14px 6px 8px; }
.pipe .n { background:var(--moss); color:#fff; border-radius:999px; width:22px;
        height:22px; display:flex; align-items:center; justify-content:center;
        font-size:12px; font-weight:700; }
.pipe .arrow { color:#9aa39a; }
.stagecard { border:1px solid #e2e6e0 !important; border-radius:14px !important;
        background:#fff !important; padding:6px 10px 10px 10px !important; }
.diffbox { line-height:1.7; padding:10px 12px; background:#fbfcfa;
        border:1px dashed #dfe5dd; border-radius:10px; }
.difflabel { text-transform:uppercase; font-size:10px; letter-spacing:0.12em;
        color:#7a847a; margin-right:8px; }
.diffbox del { background:#F7E6E4; color:#8A2F26; text-decoration:line-through;
        padding:0 3px; border-radius:4px; }
.diffbox ins { background:#E4F2E4; color:#1F5E2E; text-decoration:none;
        padding:0 3px; border-radius:4px; }
#meter { border-top:1px solid #e2e6e0; padding-top:10px; color:#4c564c; }
#truth { font-size:12px; color:#7a847a; }
"""

HEAD = """
<link rel="preconnect" href="https://fonts.googleapis.com">
<link href="https://fonts.googleapis.com/css2?family=Fraunces:opsz,wght@9..144,600&display=swap" rel="stylesheet">
"""

theme = gr.themes.Base(
    primary_hue=gr.themes.colors.green,
    neutral_hue=gr.themes.colors.stone,
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
)

with gr.Blocks(theme=theme, css=CSS, head=HEAD, title="AI Nexus · P1 demo") as demo:
    with gr.Column(elem_id="hero"):
        gr.HTML(
            '<div class="brand">AI Nexus · Green AI as a Service</div>'
            "<h1>P1 — speech that stays yours</h1>"
            '<p class="sub">Degraded speech in, clear speech out. '
            "Watch each stage of the pipeline do its part.</p>"
            '<div class="pipe">'
            '<span class="step"><span class="n">1</span>Heard · Whisper</span>'
            '<span class="arrow">→</span>'
            '<span class="step"><span class="n">2</span>Understood · small LLM</span>'
            '<span class="arrow">→</span>'
            '<span class="step"><span class="n">3</span>Spoken · Piper</span>'
            "</div>"
        )

    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            audio_in = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="Speak a sentence or two",
            )
            with gr.Accordion(
                "Personal phrase bank · personalization preview", open=False
            ):
                bank = gr.Textbox(
                    label="Words this person often uses",
                    placeholder="e.g. Collins, Hyde, physio appointment, my daughter Amara",
                    lines=2,
                )
                gr.Markdown(
                    "In the product this is a per-user adapter trained from "
                    "enrolled phrases. Here it gently biases recognition, so "
                    "you can feel what personalization does."
                )
            go = gr.Button("Run the pipeline", variant="primary")
            gr.Markdown(
                "Nothing is stored. This public demo runs in the cloud for "
                "convenience; the product runs this exact pipeline fully "
                "on-device.",
                elem_id="truth",
            )

        with gr.Column(scale=7):
            with gr.Group(elem_classes="stagecard"):
                raw_tb = gr.Textbox(label="1 · Heard — raw recognition", lines=2)
                conf_md = gr.Markdown("")
            with gr.Group(elem_classes="stagecard"):
                fixed_tb = gr.Textbox(label="2 · Understood — intended message", lines=2)
                diff_out = gr.HTML("")
            with gr.Group(elem_classes="stagecard"):
                tts_out = gr.Audio(label="3 · Spoken — their clear voice", type="filepath")

    meter_md = gr.Markdown("", elem_id="meter")

    go.click(
        run,
        inputs=[audio_in, bank],
        outputs=[raw_tb, conf_md, fixed_tb, diff_out, tts_out, meter_md],
    )

if __name__ == "__main__":
    demo.launch()