Update app.py

app.py

@@ -1,17 +1,27 @@
 """
-PlotWeaver Voice Agent — HuggingFace Space (Gradio 6 + Python 3.13)
-====================================================================
+PlotWeaver Voice Agent — HuggingFace Space (Gradio 6 + ZeroGPU)
+================================================================
 Hausa-first conversational AI for African banks, telecoms, and delivery.
-
-Pipeline (all real, running on CPU):
-  ASR  (openai/whisper-small)
-    → NLU (rule-based + Qwen2.5-1.5B-Instruct fallback, see nlu.py)
-    → Dialogue FSM (see dialogue.py)
-    → TTS  (facebook/mms-tts-hau)
-
-First turn: ~30-60s model downloads. Subsequent turns: ~5-10s on CPU.
+Pipeline:
+  ASR  (openai/whisper-small)            ── on GPU (@spaces.GPU)
+    → NLU (embedding similarity, nlu.py) ── on CPU (~200ms; GPU would be slower)
+    → Dialogue FSM (dialogue.py) + Backend (backend.py)
+    → TTS  (facebook/mms-tts-hau)        ── on GPU (@spaces.GPU)
+
+ZeroGPU notes
+-------------
+* Requires a PRO account + ZeroGPU hardware selected in Space settings.
+* Models are placed on cuda at MODULE level (the ZeroGPU contract), even though
+  a real GPU is only attached inside @spaces.GPU functions.
+* Decorated functions return CPU/numpy data (never live CUDA tensors).
+* Low `duration` values raise queue priority — ASR/TTS per turn are seconds.
+
+Secrets / variables (Space → Settings):
+  HF_TOKEN         write token, for log syncing (see logging_util.py)
+  PW_LOG_DATASET   e.g. "plotweaver/voice-agent-logs"
 """
 from __future__ import annotations
+import os
 import time
 import uuid
 import html as html_lib
@@ -20,6 +30,7 @@ from typing import Optional
 import gradio as gr
 import numpy as np
 import torch
+import spaces
 from transformers import (
     VitsModel, AutoTokenizer,
     WhisperProcessor, WhisperForConditionalGeneration,
@@ -27,42 +38,46 @@ from transformers import (
 
 from dialogue import (
     DialogueState, SCENARIOS,
-    get_prompt, get_expected_slot, transition,
+    render_prompt, get_expected_slot, transition,
 )
 from nlu import parse as nlu_parse
+from logging_util import log_turn
+
+
+# ---------------------------------------------------------------------------
+# Device selection — cuda at module level on ZeroGPU, else local CPU/GPU
+# ---------------------------------------------------------------------------
+_IS_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
+if _IS_ZEROGPU or torch.cuda.is_available():
+    DEVICE = "cuda"
+else:
+    DEVICE = "cpu"
 
 
 # ---------------------------------------------------------------------------
-# Model loading (lazy, cached)
+# Model loading (module level — required by ZeroGPU)
 # ---------------------------------------------------------------------------
-_asr_model = None
-_asr_processor = None
-_tts_model = None
-_tts_tokenizer = None
-
-
-def load_asr():
-    global _asr_model, _asr_processor
-    if _asr_model is None:
-        print("Loading Whisper-small…")
-        _asr_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-        _asr_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
-        _asr_model.eval()
-        print("Whisper-small ready.")
-    return _asr_model, _asr_processor
-
-
-def load_tts():
-    global _tts_model, _tts_tokenizer
-    if _tts_model is None:
-        print("Loading MMS-TTS Hausa…")
-        _tts_model = VitsModel.from_pretrained("facebook/mms-tts-hau")
-        _tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-hau")
-        _tts_model.eval()
-        print("MMS-TTS Hausa ready.")
-    return _tts_model, _tts_tokenizer
+print(f"Loading models on {DEVICE}…")
+
+_asr_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+_asr_model = (
+    WhisperForConditionalGeneration
+    .from_pretrained("openai/whisper-small")
+    .to(DEVICE)
+    .eval()
+)
+
+_tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-hau")
+_tts_model = VitsModel.from_pretrained("facebook/mms-tts-hau").to(DEVICE).eval()
+# PlotWeaver Hausa production voice config (tune here, not at call sites).
+_tts_model.speaking_rate = 0.85
+_tts_model.noise_scale = 0.4
+_TTS_SEED = 42
 
+print("Models ready.")
 
+
+@spaces.GPU(duration=15)
 def transcribe_hausa(audio_tuple) -> str:
     if audio_tuple is None:
         return ""
@@ -82,24 +97,26 @@ def transcribe_hausa(audio_tuple) -> str:
         num_samples = int(len(audio_array) * 16000 / sample_rate)
         audio_array = scipy.signal.resample(audio_array, num_samples).astype(np.float32)
 
-    model, processor = load_asr()
-    inputs = processor(audio_array, sampling_rate=16000, return_tensors="pt")
-    forced_ids = processor.get_decoder_prompt_ids(language="hausa", task="transcribe")
+    inputs = _asr_processor(audio_array, sampling_rate=16000, return_tensors="pt")
+    forced_ids = _asr_processor.get_decoder_prompt_ids(language="hausa", task="transcribe")
+    feats = inputs.input_features.to(DEVICE)
     with torch.no_grad():
-        ids = model.generate(inputs.input_features, forced_decoder_ids=forced_ids, max_new_tokens=128)
-    text = processor.batch_decode(ids, skip_special_tokens=True)[0].strip()
+        ids = _asr_model.generate(feats, forced_decoder_ids=forced_ids, max_new_tokens=128)
+    text = _asr_processor.batch_decode(ids, skip_special_tokens=True)[0].strip()
     return text
 
 
+@spaces.GPU(duration=10)
 def synthesize_hausa(text: str) -> Optional[tuple]:
     if not text.strip():
         return None
-    model, tokenizer = load_tts()
-    inputs = tokenizer(text, return_tensors="pt")
+    torch.manual_seed(_TTS_SEED)  # reproducible voice (prod parity)
+    inputs = _tts_tokenizer(text, return_tensors="pt")
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
     with torch.no_grad():
-        out = model(**inputs).waveform
+        out = _tts_model(**inputs).waveform
     audio = out.squeeze().cpu().numpy().astype(np.float32)
-    return (model.config.sampling_rate, audio)
+    return (_tts_model.config.sampling_rate, audio)
 
 
 # ---------------------------------------------------------------------------
@@ -198,15 +215,17 @@ def render_whatsapp(session: dict) -> str:
 # ---------------------------------------------------------------------------
 def run_turn(user_text: str, session: dict, is_voice: bool = False):
     """Returns (updated_session_dict, bot_audio)."""
+    t0 = time.time()
     state = DialogueState.from_dict(session) if session else None
     if state is None:
         state = DialogueState(session_id="sess_" + uuid.uuid4().hex[:8], vertical="bank")
 
+    state_from = state.current_state
     expected = get_expected_slot(state.vertical, state.current_state)
-    intent, entities, _ = nlu_parse(user_text, expected)
+    intent, entities, source = nlu_parse(user_text, expected)
     state = transition(state, intent, entities)
 
-    prompt = get_prompt(state.vertical, state.current_state)
+    prompt = render_prompt(state.vertical, state.current_state, state.slots)
 
     state.history.append({"role": "user", "text": user_text, "is_voice": is_voice})
     state.history.append({"role": "bot", "text_ha": prompt["ha"], "text_en": prompt["en"]})
@@ -217,6 +236,21 @@ def run_turn(user_text: str, session: dict, is_voice: bool = False):
         print(f"TTS failed: {e}")
         audio = None
 
+    log_turn({
+        "session_id": state.session_id,
+        "vertical": state.vertical,
+        "turn": state.turn_count,
+        "is_voice": is_voice,
+        "asr_text": user_text,
+        "expected_slot": expected,
+        "intent": intent,
+        "nlu_source": source,
+        "state_from": state_from,
+        "state_to": state.current_state,
+        "escalated": state.escalate_to_human,
+        "latency_ms": round((time.time() - t0) * 1000),
+    })
+
     return state.to_dict(), audio
 
 
@@ -225,7 +259,7 @@ def run_turn(user_text: str, session: dict, is_voice: bool = False):
 # ---------------------------------------------------------------------------
 def on_vertical_change(vertical: str):
     state = DialogueState(session_id="sess_" + uuid.uuid4().hex[:8], vertical=vertical)
-    greet = get_prompt(vertical, "greeting")
+    greet = render_prompt(vertical, "greeting", state.slots)
     state.history.append({"role": "bot", "text_ha": greet["ha"], "text_en": greet["en"]})
     session = state.to_dict()
     return session, render_whatsapp(session), None
@@ -341,4 +375,4 @@ with gr.Blocks(css=CUSTOM_CSS, title="PlotWeaver Voice Agent") as demo:
 
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)