Update app.py

app.py

@@ -1,13 +1,16 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import FileResponse
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import whisper
 import torch
+from gtts import gTTS
 import os
 
 app = FastAPI()
-model_name = "Qwen/Qwen3-4B-Instruct-2507"
 
-# Load tokenizer và model (CPU cho Spaces Free)
+# Load Qwen model
+model_name = "Qwen/Qwen3-4B-Instruct-2507"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -15,7 +18,11 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float32
 )
 
-conversation = [{"role": "system", "content": "Bạn là một trợ lý AI. Hãy trả lời ngắn gọn, súc tích, tối đa 2 câu."}]  # Lưu hội thoại
+# Load Whisper model
+whisper_model = whisper.load_model("base")
+
+# Lưu hội thoại
+conversation = [{"role": "system", "content": "Bạn là một trợ lý AI. Hãy trả lời ngắn gọn, súc tích, tối đa 2 câu."}]
 
 class ChatRequest(BaseModel):
     message: str
@@ -24,25 +31,55 @@ class ChatRequest(BaseModel):
 def read_root():
     return {"message": "Ứng dụng đang chạy!"}
 
+# Endpoint chat text
 @app.post("/chat")
 async def chat(request: ChatRequest):
     conversation.append({"role": "user", "content": request.message})
-
-    # Áp dụng template hội thoại
     text = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    response_text = generate_full_response(model_inputs)
+    conversation.append({"role": "assistant", "content": response_text})
+    return {"response": response_text}
+
+# Endpoint voice chat + TTS
+@app.post("/voice_chat")
+async def voice_chat(file: UploadFile = File(...)):
+    # Lưu file tạm
+    file_location = f"temp_{file.filename}"
+    with open(file_location, "wb") as f:
+        f.write(await file.read())
 
+    # Chuyển âm thanh thành text
+    result = whisper_model.transcribe(file_location, language="vi")
+    user_text = result["text"]
+
+    # Gọi mô hình Qwen để trả lời
+    conversation.append({"role": "user", "content": user_text})
+    text = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
     response_text = generate_full_response(model_inputs)
     conversation.append({"role": "assistant", "content": response_text})
 
-    return {"response": response_text}
+    # Tạo file âm thanh từ phản hồi
+    tts = gTTS(response_text, lang="vi")
+    audio_file = "response.mp3"
+    tts.save(audio_file)
 
+    return {
+        "user_text": user_text,
+        "response": response_text,
+        "audio_url": f"/get_audio"
+    }
 
+# Endpoint trả về file âm thanh
+@app.get("/get_audio")
+async def get_audio():
+    return FileResponse("response.mp3", media_type="audio/mpeg")
+
+# Hàm sinh phản hồi
 def generate_full_response(model_inputs, max_new_tokens=64):
     with torch.inference_mode():
         generated_ids = model.generate(**model_inputs, max_new_tokens=max_new_tokens)
     output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
     response_text = tokenizer.decode(output_ids, skip_special_tokens=True)
-
-
     return response_text.strip()