app.py

import spaces  # Import the ZeroGPU helper
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch
from torch.nn.functional import softmax
import numpy as np
import soundfile as sf
import io
import tempfile
import outlines  # For Qwen integration via outlines
import kokoro     # For TTS synthesis
import re
from pathlib import Path
from functools import lru_cache
import warnings

# Suppress FutureWarnings (e.g. about using `inputs` vs. `input_features`)
warnings.filterwarnings("ignore", category=FutureWarning)

# ------------------- Model Identifiers -------------------
whisper_model_id = "Jingmiao/whisper-small-zh_tw"
qwen_model_id = "Qwen/Qwen2.5-0.5B-Instruct"

available_models = {
    "ALBERT-tiny (Chinese)": "Luigi/albert-tiny-chinese-dinercall-intent",
    "ALBERT-base (Chinese)": "Luigi/albert-base-chinese-dinercall-intent",
    "Qwen (via Transformers - outlines)": "qwen"
}

# ------------------- Caching and Loading Functions -------------------
@lru_cache(maxsize=1)
def load_whisper_pipeline():
    pipe = pipeline("automatic-speech-recognition", 
                    model=whisper_model_id,
                    chunk_length_s=30)
    # Move model to GPU if available for faster inference
    if torch.cuda.is_available():
        pipe.model.to("cuda")
    return pipe

@lru_cache(maxsize=2)
def load_transformers_model(model_id: str):
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    model = AutoModelForSequenceClassification.from_pretrained(model_id)
    if torch.cuda.is_available():
        model.to("cuda")
    return tokenizer, model

@lru_cache(maxsize=1)
def load_qwen_model():
    return outlines.models.transformers(qwen_model_id)

@lru_cache(maxsize=1)
def get_tts_pipeline():
    return kokoro.KPipeline(lang_code="z")

# ------------------- Inference Functions -------------------
def predict_with_qwen(text: str):
    model = load_qwen_model()
    prompt = f"""
<|im_start|>system
You are an expert in classification of restaurant customers' messages.
You must decide between the following two intents:
RESERVATION: Inquiries and requests highly related to table reservations and seating.
NOT_RESERVATION: All other messages.
Respond with *only* the intent label in a JSON object, like: {{"result": "RESERVATION"}}.
<|im_end|>

<|im_start|>user
Classify the following message: "{text}"
<|im_end|>

<|im_start|>assistant
"""
    generator = outlines.generate.choice(model, ["RESERVATION", "NOT_RESERVATION"])
    prediction = generator(prompt)
    if prediction == "RESERVATION":
        return "📞 訂位意圖 (Reservation intent)"
    elif prediction == "NOT_RESERVATION":
        return "❌ 無訂位意圖 (Not Reservation intent)"
    else:
        return f"未知回應: {prediction}"

def predict_intent(text: str, model_id: str):
    tokenizer, model = load_transformers_model(model_id)
    inputs = tokenizer(text, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.to("cuda") for k, v in inputs.items()}
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = softmax(logits, dim=-1)
        confidence = probs[0, 1].item()
    if confidence >= 0.7:
        return f"📞 訂位意圖 (Reservation intent)（訂位信心度 Confidence: {confidence:.2%}）"
    else:
        return f"❌ 無訂位意圖 (Not Reservation intent)（訂位信心度 Confidence: {confidence:.2%}）"

def get_tts_message(intent_result: str):
    if intent_result and "訂位意圖" in intent_result and "無" not in intent_result:
        return "稍後您將會從簡訊收到訂位連結"
    elif intent_result:
        return "我們將會將您的回饋傳達給負責人，謝謝您"
    else:
        return "未能判斷意圖"

def tts_audio_output(message: str, voice: str = 'af_heart'):
    pipeline_tts = get_tts_pipeline()
    generator = pipeline_tts(message, voice=voice)
    audio_chunks = []
    for _, _, audio in generator:
        audio_chunks.append(audio)
    if audio_chunks:
        audio_concat = np.concatenate(audio_chunks)
        # Return as tuple (sample_rate, numpy_array) for gr.Audio (using 24000 Hz)
        return (24000, audio_concat)
    else:
        return None

def transcribe_audio(audio_input):
    whisper_pipe = load_whisper_pipeline()
    # For file input, audio_input is a filepath string.
    if isinstance(audio_input, str):
        result = whisper_pipe(audio_input)
        return result["text"]
    # For microphone input, we now also use file_path.
    elif isinstance(audio_input, tuple):
        # In our updated configuration, microphone input should be provided as a file path,
        # so this branch may not be reached.
        return ""
    else:
        return ""

# ------------------- Main Processing Function -------------------
@spaces.GPU  # Decorate to run on GPU when processing
def classify_intent(mode, mic_audio, text_input, file_audio, model_choice):
    # Determine input based on selected mode.
    if mode == "Microphone" and mic_audio is not None:
        # mic_audio is a file path.
        transcription = transcribe_audio(mic_audio)
    elif mode == "Text" and text_input:
        transcription = text_input
    elif mode == "File" and file_audio is not None:
        transcription = transcribe_audio(file_audio)
    else:
        return "請提供語音或文字輸入", "", None

    # Classify the transcribed or provided text.
    if available_models[model_choice] == "qwen":
        classification = predict_with_qwen(transcription)
    else:
        classification = predict_intent(transcription, available_models[model_choice])
    # Generate TTS message and corresponding audio.
    tts_msg = get_tts_message(classification)
    tts_audio = tts_audio_output(tts_msg)
    return transcription, classification, tts_audio

# ------------------- Gradio Blocks Interface Setup -------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🍽️ 餐廳訂位意圖識別")
    gr.Markdown("錄音、上傳語音檔案或輸入文字，自動判斷是否具有訂位意圖。")
    
    with gr.Row():
        mode = gr.Radio(choices=["Microphone", "Text", "File"], label="選擇輸入模式", value="Microphone")
    
    with gr.Row():
        # For microphone input, set type="filepath" so that we always get a file path.
        mic_audio = gr.Audio(sources=["microphone"], type="filepath", label="語音輸入 (點擊錄音)")
        text_input = gr.Textbox(lines=2, placeholder="請輸入文字", label="文字輸入")
        file_audio = gr.Audio(sources=["upload"], type="filepath", label="上傳語音檔案")
    
    # Initially, only the microphone input is visible.
    text_input.visible = False
    file_audio.visible = False

    # Set visibility based on selected mode.
    def update_visibility(selected_mode):
        if selected_mode == "Microphone":
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
        elif selected_mode == "Text":
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
        else:  # File
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
    mode.change(fn=update_visibility, inputs=mode, outputs=[mic_audio, text_input, file_audio])
    
    with gr.Row():
        model_dropdown = gr.Dropdown(choices=list(available_models.keys()),
                                     value="ALBERT-tiny (Chinese)", label="選擇模型")
    
    with gr.Row():
        classify_btn = gr.Button("執行辨識")
    
    with gr.Row():
        transcription_output = gr.Textbox(label="轉換文字")
    with gr.Row():
        classification_output = gr.Textbox(label="意圖判斷結果")
    with gr.Row():
        tts_output = gr.Audio(type="numpy", label="TTS 語音輸出")
    
    # Button event triggers the classification.
    classify_btn.click(fn=classify_intent, 
                       inputs=[mode, mic_audio, text_input, file_audio, model_dropdown],
                       outputs=[transcription_output, classification_output, tts_output])

demo.launch()