from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.responses import HTMLResponse
from huggingface_hub import InferenceClient
import os
import uvicorn
import base64
import requests
# Khởi tạo FastAPI
app = FastAPI()
# Lấy API key từ biến môi trường
HF_API_KEY = os.getenv("HF_API_KEY")
if not HF_API_KEY:
raise ValueError("❌ Thiếu HF_API_KEY! Thêm vào biến môi trường.")
# Khởi tạo Hugging Face Client
client = InferenceClient(token=HF_API_KEY)
# Tạo thư mục lưu file tạm
TEMP_DIR = "temp"
os.makedirs(TEMP_DIR, exist_ok=True)
# Mô hình sử dụng
STT_MODEL = "openai/whisper-tiny.en"
TTS_MODEL = "facebook/mms-tts-eng"
CHAT_MODEL = "mistralai/Mistral-7B-Instruct-v0.1" # Model nhẹ hơn
# Giao diện HTML
HTML_CONTENT = """
Chatbot AI
Chatbot AI
"""
@app.get("/", response_class=HTMLResponse)
async def root():
return HTML_CONTENT
async def text_to_speech(text: str) -> str:
output_path = os.path.join(TEMP_DIR, "output.wav")
try:
print(f"🔄 Đang xử lý TTS với model {TTS_MODEL}...")
response = client.post(f"{TTS_MODEL}", json={"inputs": text})
with open(output_path, "wb") as f:
f.write(response)
with open(output_path, "rb") as f:
audio_data = f.read()
return base64.b64encode(audio_data).decode('utf-8')
except Exception as e:
print(f"❌ Lỗi TTS: {e}")
raise
async def speech_to_text(file: UploadFile) -> str:
try:
print(f"🔄 Đang xử lý STT với model {STT_MODEL}...")
audio_data = await file.read()
response = client.post(f"{STT_MODEL}", files={"file": audio_data})
return response.get("text", "Không nghe được gì.")
except Exception as e:
print(f"❌ Lỗi STT: {e}")
raise
async def generate_text(prompt: str) -> str:
try:
print(f"🔄 Đang xử lý hội thoại với model {CHAT_MODEL}...")
response = client.post(f"{CHAT_MODEL}", json={"inputs": prompt})
return response.get("generated_text", "Xin lỗi, tôi gặp lỗi khi xử lý câu hỏi của bạn.").strip()
except Exception as e:
print(f"❌ Lỗi xử lý hội thoại: {e}")
return "Xin lỗi, tôi gặp lỗi khi xử lý câu hỏi của bạn."
@app.post("/chatbot")
async def chatbot(prompt: str = Form(None), file: UploadFile = None):
try:
if file:
stt_output = await speech_to_text(file)
chat_output = await generate_text(stt_output)
elif prompt:
chat_output = await generate_text(prompt)
else:
raise HTTPException(status_code=422, detail="Thiếu dữ liệu đầu vào")
audio_base64 = await text_to_speech(chat_output)
return {"text": chat_output, "audio": audio_base64}
except HTTPException as he:
raise he
except Exception as e:
return {"error": str(e)}
if __name__ == "__main__":
print("🚀 Khởi động FastAPI Server...")
uvicorn.run(app, host="0.0.0.0", port=7860)