Update app.py

app.py

@@ -1,93 +1,110 @@
 import gradio as gr
-import os
 import torch
 import json
-import base64
 import tempfile
-
+import faiss
 from gtts import gTTS
-import whisper
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
-import faiss
-
-# ---------- 模型與資料載入 ---------- #
-
-# 問答模型（輕量中文 LLM）
-LLM_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
-model = AutoModelForCausalLM.from_pretrained(LLM_MODEL).eval()
-
-# 向量模型 + 向量資料庫
-embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-index = faiss.read_index("vector_store.index")
-
-# 讀取文本資料（編號與原始句）
-with open("chunks.json", "r", encoding="utf-8") as f:
-    chunks = json.load(f)
-
-# 語音辨識模型（Whisper）
-asr_model = whisper.load_model("base")
-
-# ---------- 問答處理 ---------- #
-
-def generate_answer(query):
-    embedding = embedder.encode([query])
-    D, I = index.search(embedding, k=3)
-    context = "\n".join([chunks[i] for i in I[0]])
-
-    prompt = f"你是一位語音問答助手，請根據下方資訊回答問題。\n\n資訊：\n{context}\n\n問題：{query}\n\n回答："
-
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=128)
-    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+import numpy as np
+
+# 模型
+MODEL_NAME = "openbmb/MiniCPM-2B-sft-bf16"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True).eval()
+
+# 語音辨識 Whisper
+asr = pipeline("automatic-speech-recognition", model="openai/whisper-small", device=0 if torch.cuda.is_available() else -1)
+
+# 向量模型
+encoder = SentenceTransformer("shibing624/text2vec-base-chinese")
+index = faiss.read_index("vector_store.faiss")
+with open("documents.json", "r", encoding="utf-8") as f:
+    documents = json.load(f)
+
+# QA固定問答（可選）
+try:
+    with open("qa.json", "r", encoding="utf-8") as f:
+        qa_data = json.load(f)
+except:
+    qa_data = []
+
+# QA match（選擇性）
+def match_qa(user_input):
+    cleaned_input = user_input.replace(" ", "")
+    for item in qa_data:
+        if item["match"] == "OR":
+            if any(k.replace(" ", "") in cleaned_input for k in item["keywords"]):
+                return item["response"]
+        elif item["match"] == "AND":
+            if all(k.replace(" ", "") in cleaned_input for k in item["keywords"]):
+                return item["response"]
+    return None
+
+# 文字生成
+def generate_answer(text):
+    messages = [{"role": "user", "content": text}]
+    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(input_ids, max_new_tokens=200)
+    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
     return response.strip()
 
-# ---------- 語音處理 ---------- #
-
-def asr(audio_path):
-    result = asr_model.transcribe(audio_path, language="zh")
-    return result["text"]
-
-def tts(text):
-    tts = gTTS(text, lang="zh")
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as fp:
-        tts.save(fp.name)
-        return fp.name
-
-# ---------- Pipeline ---------- #
-
-def chat_pipeline(audio_input=None, text_input=None):
+# 向量比對
+def search_vector_db(query, top_k=1):
+    q_vec = encoder.encode([query])
+    D, I = index.search(np.array(q_vec), top_k)
+    results = [documents[i] for i in I[0] if i < len(documents)]
+    return results
+
+# 回答邏輯整合
+def answer(text):
+    # 1. QA 固定資料庫
+    fixed = match_qa(text)
+    if fixed:
+        return fixed
+
+    # 2. RAG 取資料輔助
+    related_docs = search_vector_db(text)
+    context = "\n".join(related_docs)
+    prompt = f"以下是一些關於南臺科技大學的資料：\n{context}\n\n根據上面的資料，請用中文簡短回答這個問題：{text}"
+    return generate_answer(prompt)
+
+# TTS
+def text_to_speech(text):
+    tts = gTTS(text, lang='zh')
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
+    tts.save(tmp.name)
+    return tmp.name
+
+# 主流程
+def voice_assistant(audio_input=None, text_input=None):
     if audio_input:
-        text = asr(audio_input)
+        result = asr(audio_input)
+        user_text = result["text"]
     elif text_input:
-        text = text_input
+        user_text = text_input
     else:
-        return "請輸入問題或語音", None
-
-    answer = generate_answer(text)
-    audio_out = tts(answer)
-    return answer, audio_out
+        return "請輸入語音或文字", None
 
-# ---------- Gradio 介面 ---------- #
+    response = answer(user_text)
+    speech_file = text_to_speech(response)
+    return response, speech_file
 
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎙️ 南臺科技大學 問答語音助理（TinyLlama + Whisper + RAG）")
+    gr.Markdown("## 🎓 南臺科技大學 AI 語音助理（MiniCPM + Whisper + 向量式 RAG）")
 
     with gr.Row():
-        audio_input = gr.Audio(source="microphone", type="filepath", label="🎤 語音提問")
-        text_input = gr.Textbox(label="或文字輸入", placeholder="請輸入您的問題")
+        mic = gr.Audio(source="microphone", type="filepath", label="語音輸入")
+        text_input = gr.Textbox(label="文字輸入", placeholder="請輸入您的問題")
 
     submit_btn = gr.Button("送出")
 
-    output_text = gr.Textbox(label="AI 回答")
+    output_text = gr.Textbox(label="回答")
     output_audio = gr.Audio(label="語音播放", type="filepath")
 
-    submit_btn.click(
-        fn=chat_pipeline,
-        inputs=[audio_input, text_input],
-        outputs=[output_text, output_audio]
-    )
+    submit_btn.click(fn=voice_assistant, inputs=[mic, text_input], outputs=[output_text, output_audio])
 
 if __name__ == "__main__":
     demo.launch()