Minhtan210905/Toan-Roi-Rac-87M
0
1import torch2import torch.nn.functional as F3import pickle4import os5import re6import sys7import csv8import time9import json as json_module10import uuid11import asyncio12from typing import List, Optional13from fastapi import FastAPI, HTTPException, Path14from fastapi.staticfiles import StaticFiles15from fastapi.responses import FileResponse, StreamingResponse16from fastapi.middleware.cors import CORSMiddleware17from pydantic import BaseModel18from contextlib import asynccontextmanager19import sqlite320 21if sys.stdout.encoding != 'utf-8':22 sys.stdout.reconfigure(encoding='utf-8')23 24# ── Paths ──25VOCAB_PATH = 'data/processed/vocab.pkl'26MODEL_PATH = 'models/best/model.pt'27TOKENIZER_PATH = 'data/processed/bpe_tokenizer.json'28RAW_CSV_PATH = 'data/raw/output.csv'29DB_PATH = 'data/chatbot.sqlite'30 31# ── Model globals ──32model = None33word2idx = None34idx2word = None35device = None36model_info = {}37bpe_tokenizer = None38 39qa_database = []40qa_embeddings = None41SIMILARITY_THRESHOLD = 0.9242 43 44# ═══════════════════════════════════════45# DATABASE – Sessions + Messages46# ═══════════════════════════════════════47 48def init_db():49 conn = sqlite3.connect(DB_PATH)50 c = conn.cursor()51 c.execute("PRAGMA foreign_keys = ON")52 c.execute('''53 CREATE TABLE IF NOT EXISTS sessions (54 id TEXT PRIMARY KEY,55 title TEXT NOT NULL DEFAULT 'Chat mới',56 created_at DATETIME DEFAULT CURRENT_TIMESTAMP,57 updated_at DATETIME DEFAULT CURRENT_TIMESTAMP58 )59 ''')60 c.execute('''61 CREATE TABLE IF NOT EXISTS messages (62 id INTEGER PRIMARY KEY AUTOINCREMENT,63 session_id TEXT NOT NULL,64 role TEXT NOT NULL,65 content TEXT NOT NULL,66 timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,67 FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE68 )69 ''')70 # Migrate old chat_history table if it exists71 c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='chat_history'")72 if c.fetchone():73 c.execute("SELECT COUNT(*) FROM chat_history")74 count = c.fetchone()[0]75 if count > 0:76 sid = str(uuid.uuid4())77 c.execute("INSERT INTO sessions (id, title) VALUES (?, ?)", (sid, "Chat cũ"))78 c.execute("""79 INSERT INTO messages (session_id, role, content, timestamp)80 SELECT ?, role, content, timestamp FROM chat_history81 """, (sid,))82 print(f"[DB] Migrated {count} old messages to session '{sid}'")83 c.execute("DROP TABLE chat_history")84 conn.commit()85 conn.close()86 87 88def save_message(session_id, role, content):89 conn = sqlite3.connect(DB_PATH)90 c = conn.cursor()91 c.execute("INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)",92 (session_id, role, content))93 c.execute("UPDATE sessions SET updated_at = CURRENT_TIMESTAMP WHERE id = ?",94 (session_id,))95 conn.commit()96 conn.close()97 98 99def create_session(title="Chat mới"):100 sid = str(uuid.uuid4())101 conn = sqlite3.connect(DB_PATH)102 c = conn.cursor()103 c.execute("INSERT INTO sessions (id, title) VALUES (?, ?)", (sid, title))104 conn.commit()105 conn.close()106 return sid107 108 109def update_session_title(session_id, title):110 conn = sqlite3.connect(DB_PATH)111 c = conn.cursor()112 c.execute("UPDATE sessions SET title = ? WHERE id = ?", (title, session_id))113 conn.commit()114 conn.close()115 116 117# ═══════════════════════════════════════118# MODEL LOADING119# ═══════════════════════════════════════120 121def load_bpe_tokenizer():122 global bpe_tokenizer123 if os.path.exists(TOKENIZER_PATH):124 from tokenizers import Tokenizer125 bpe_tokenizer = Tokenizer.from_file(TOKENIZER_PATH)126 print(f"[OK] BPE Tokenizer loaded | Vocab: {bpe_tokenizer.get_vocab_size()}")127 return True128 else:129 print(f"[WARN] Không tìm thấy BPE tokenizer: {TOKENIZER_PATH}")130 return False131 132 133def tokenize_bpe(text):134 if bpe_tokenizer is not None:135 enc = bpe_tokenizer.encode(str(text).strip())136 return enc.ids137 else:138 text = str(text).lower()139 text = re.sub(r'([.,!?()])', r' \\1 ', text)140 text = re.sub(r'\\s{2,}', ' ', text)141 tokens = text.strip().split()142 return [word2idx.get(w, word2idx.get('<unk>', 3)) for w in tokens]143 144 145def clean_response(text):146 text = text.strip()147 if not text:148 return ""149 text = text[0].upper() + text[1:]150 import unicodedata151 text = unicodedata.normalize('NFC', text)152 text = re.sub(r'\\\s+([a-zA-Z]+)', r'\\\1', text)153 text = re.sub(r'\s*_\s*', '_', text)154 text = re.sub(r'\s*\^\s*', '^', text)155 text = re.sub(r'\s*\{\s*', '{', text)156 text = re.sub(r'\s*\}\s*', '}', text)157 text = re.sub(r'\s+([.,!?():;])', r'\1', text)158 text = re.sub(r'\(\s+', r'(', text)159 text = re.sub(r'\s{2,}', ' ', text)160 return text.strip()161 162 163def generate_response(message, temperature=0.5, top_k=3, top_p=0.9,164 repetition_penalty=1.3, max_tokens=100):165 """Core inference function. Returns (response_text, tokens_count, elapsed_ms)."""166 start = time.perf_counter()167 168 # Greeting check169 msg_lower = message.lower()170 if len(msg_lower) < 25:171 if any(w in msg_lower for w in ["chào", "hello", "hi ", "xin chào", "alo"]):172 return "Chào bạn! Mình là trợ lý môn Toán Rời Rạc. Bạn cần hỏi gì nào?", 0, 1.0173 if any(w in msg_lower for w in ["cảm ơn", "thank", "cám ơn", "tks"]):174 return "Không có gì! Chúc bạn học tốt môn Toán Rời Rạc nhé!", 0, 1.0175 176 tokens = tokenize_bpe(message)177 if not tokens:178 return "Mình không hiểu câu hỏi. Bạn thử hỏi lại nhé!", 0, 1.0179 180 sos_id = word2idx.get('<sos>', 1)181 sep_id = word2idx.get('<sep>', 4)182 183 prompt = [sos_id] + tokens + [sep_id]184 x = torch.tensor([prompt]).to(device)185 186 with torch.no_grad():187 out_ids = model.generate(x, max_new_tokens=max_tokens,188 temperature=temperature,189 top_k=top_k, top_p=top_p,190 repetition_penalty=repetition_penalty)191 192 elapsed_ms = (time.perf_counter() - start) * 1000193 resp_ids = out_ids[0][len(prompt):]194 195 if bpe_tokenizer is not None:196 valid_ids = []197 for i in resp_ids:198 tid = i.item()199 if tid == word2idx.get('<eos>', 2):200 break201 if tid not in [word2idx.get('<pad>', 0), word2idx.get('<sos>', 1),202 word2idx.get('<unk>', 3), word2idx.get('<sep>', 4)]:203 valid_ids.append(tid)204 response = bpe_tokenizer.decode(valid_ids)205 else:206 resp_tokens = []207 for i in resp_ids:208 sym = idx2word.get(i.item(), '<unk>')209 if sym == '<eos>':210 break211 if sym not in ['<pad>', '<sos>', '<unk>', '<sep>']:212 resp_tokens.append(sym)213 response = " ".join(resp_tokens)214 215 response = clean_response(response)216 if not response.strip():217 response = "Mình chưa học tới phần này. Bạn hỏi câu khác nhé!"218 219 return response, len(resp_ids), round(elapsed_ms, 2)220 221 222async def generate_response_stream(message, temperature=0.5, top_k=3, top_p=0.9,223 repetition_penalty=1.3, max_tokens=100):224 """Streaming inference function. Yields SSE data chunks containing the accumulated full text."""225 msg_lower = message.lower()226 if len(msg_lower) < 25:227 if any(w in msg_lower for w in ["chào", "hello", "hi ", "xin chào", "alo"]):228 yield "data: " + json_module.dumps({"text": "Chào bạn! Mình là trợ lý môn Toán Rời Rạc. Bạn cần hỏi gì nào?", "done": True}, ensure_ascii=False) + "\n\n"229 return230 if any(w in msg_lower for w in ["cảm ơn", "thank", "cám ơn", "tks"]):231 yield "data: " + json_module.dumps({"text": "Không có gì! Chúc bạn học tốt môn Toán Rời Rạc nhé!", "done": True}, ensure_ascii=False) + "\n\n"232 return233 234 tokens = tokenize_bpe(message)235 if not tokens:236 yield "data: " + json_module.dumps({"text": "Mình không hiểu câu hỏi. Bạn thử hỏi lại nhé!", "done": True}, ensure_ascii=False) + "\n\n"237 return238 239 sos_id = word2idx.get('<sos>', 1)240 sep_id = word2idx.get('<sep>', 4)241 eos_id = word2idx.get('<eos>', 2)242 pad_id = word2idx.get('<pad>', 0)243 unk_id = word2idx.get('<unk>', 3)244 245 prompt = [sos_id] + tokens + [sep_id]246 x = torch.tensor([prompt]).to(device)247 248 valid_ids = []249 250 with torch.no_grad():251 for tid in model.generate_stream(x, max_new_tokens=max_tokens,252 temperature=temperature, top_k=top_k,253 top_p=top_p, repetition_penalty=repetition_penalty):254 if tid == eos_id:255 break256 if tid not in [pad_id, sos_id, unk_id, sep_id]:257 valid_ids.append(tid)258 259 if bpe_tokenizer is not None:260 current_text = bpe_tokenizer.decode(valid_ids)261 else:262 current_text = " ".join([idx2word.get(i, '<unk>') for i in valid_ids])263 264 current_text = clean_response(current_text)265 yield "data: " + json_module.dumps({"text": current_text, "done": False}, ensure_ascii=False) + "\n\n"266 267 await asyncio.sleep(0.04)268 269 if not valid_ids:270 yield "data: " + json_module.dumps({"text": "Mình chưa học tới phần này. Bạn hỏi câu khác nhé!", "done": True}, ensure_ascii=False) + "\n\n"271 else:272 yield "data: " + json_module.dumps({"done": True}, ensure_ascii=False) + "\n\n"273 274 275def load_model():276 global model, word2idx, idx2word, device, model_info277 sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))278 from model import make_model279 280 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")281 282 if not os.path.exists(VOCAB_PATH):283 print(f"[WARN] Không tìm thấy vocab: {VOCAB_PATH}")284 return False285 if not os.path.exists(MODEL_PATH):286 print(f"[WARN] Không tìm thấy model: {MODEL_PATH}")287 return False288 289 with open(VOCAB_PATH, 'rb') as f:290 vocab = pickle.load(f)291 word2idx = vocab['word2idx']292 idx2word = vocab['idx2word']293 294 ckpt = torch.load(MODEL_PATH, map_location=device)295 d_m = ckpt.get('d_m', 768)296 n_l = ckpt.get('n_l', 12)297 n_h = ckpt.get('n_h', 12)298 dr = ckpt.get('dr', 0.1)299 300 model = make_model(len(word2idx), n_layer=n_l, d_model=d_m, n_head=n_h, dropout=dr)301 model.load_state_dict(ckpt['model_state_dict'])302 model.to(device)303 model.eval()304 305 model_info = {306 "vocab_size": len(word2idx),307 "d_model": d_m,308 "n_layers": n_l,309 "n_heads": n_h,310 "device": str(device),311 }312 print(f"[OK] Model loaded on {device} | d_model={d_m} | layers={n_l}")313 return True314 315 316# ═══════════════════════════════════════317# FASTAPI APP318# ═══════════════════════════════════════319 320@asynccontextmanager321async def lifespan(app: FastAPI):322 success = load_model()323 load_bpe_tokenizer()324 init_db()325 if success:326 print("[INFO] Server ready.")327 else:328 print("[WARN] Server khởi động KHÔNG có model.")329 yield330 331 332app = FastAPI(title="Toán Rời Rạc Chatbot", version="3.0.0", lifespan=lifespan)333 334app.add_middleware(335 CORSMiddleware,336 allow_origins=["*"],337 allow_credentials=True,338 allow_methods=["*"],339 allow_headers=["*"],340)341 342 343# ── Health ──344@app.get("/api/health")345async def health_check():346 return {347 "status": "ok",348 "model_loaded": model is not None,349 "model_info": model_info,350 "bpe_tokenizer": bpe_tokenizer is not None,351 }352 353 354# ── Sessions ──355@app.get("/api/sessions")356async def list_sessions():357 conn = sqlite3.connect(DB_PATH)358 conn.row_factory = sqlite3.Row359 c = conn.cursor()360 c.execute("""361 SELECT s.id, s.title, s.created_at, s.updated_at,362 (SELECT COUNT(*) FROM messages m WHERE m.session_id = s.id) as msg_count363 FROM sessions s ORDER BY s.updated_at DESC364 """)365 rows = c.fetchall()366 conn.close()367 return [dict(r) for r in rows]368 369 370@app.post("/api/sessions")371async def create_session_endpoint():372 sid = create_session()373 return {"id": sid, "title": "Chat mới"}374 375 376@app.delete("/api/sessions/{session_id}")377async def delete_session(session_id: str = Path(...)):378 conn = sqlite3.connect(DB_PATH)379 c = conn.cursor()380 c.execute("PRAGMA foreign_keys = ON")381 c.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))382 c.execute("DELETE FROM sessions WHERE id = ?", (session_id,))383 conn.commit()384 conn.close()385 return {"status": "ok"}386 387 388@app.get("/api/sessions/{session_id}/messages")389async def get_session_messages(session_id: str = Path(...)):390 conn = sqlite3.connect(DB_PATH)391 conn.row_factory = sqlite3.Row392 c = conn.cursor()393 c.execute("SELECT role, content, timestamp FROM messages WHERE session_id = ? ORDER BY id ASC",394 (session_id,))395 rows = c.fetchall()396 conn.close()397 return [dict(r) for r in rows]398 399 400# ── Streaming Chat ──401class StreamChatRequest(BaseModel):402 message: str403 session_id: Optional[str] = None404 temperature: float = 0.5405 top_k: int = 3406 top_p: float = 0.9407 repetition_penalty: float = 1.3408 max_tokens: int = 100409 410 411@app.post("/api/chat/stream")412async def chat_stream(req: StreamChatRequest):413 if model is None:414 raise HTTPException(status_code=503, detail="Model chưa được load.")415 416 message = req.message.strip()417 if not message:418 raise HTTPException(status_code=400, detail="Tin nhắn không được để trống.")419 420 # Create or use existing session421 session_id = req.session_id422 is_new_session = False423 if not session_id:424 session_id = create_session(message[:40])425 is_new_session = True426 427 # Save user message428 save_message(session_id, "user", message)429 430 # If new session, update title with first message431 if is_new_session:432 title = message[:40] + ("..." if len(message) > 40 else "")433 update_session_title(session_id, title)434 435 async def event_stream():436 # Send a 1024-byte padding comment to bypass browser buffering437 yield f": {' ' * 1024}\n\n"438 439 # Send session info440 yield f"data: {json_module.dumps({'type': 'session', 'session_id': session_id, 'is_new': is_new_session})}\n\n"441 442 # Send thinking indicator443 yield f"data: {json_module.dumps({'type': 'thinking'})}\n\n"444 await asyncio.sleep(0.05)445 446 # Generate response using real streaming447 token_count = 0448 start_time = time.perf_counter()449 final_response = ""450 451 async for chunk in generate_response_stream(452 message,453 temperature=req.temperature,454 top_k=req.top_k,455 top_p=req.top_p,456 repetition_penalty=req.repetition_penalty,457 max_tokens=req.max_tokens,458 ):459 # Parse chunk JSON460 data = json_module.loads(chunk[6:]) # strip 'data: '461 if not data.get("done"):462 final_response = data["text"]463 yield f"data: {json_module.dumps({'type': 'token', 'content': final_response})}\n\n"464 else:465 if "text" in data:466 final_response = data["text"]467 yield f"data: {json_module.dumps({'type': 'token', 'content': final_response})}\n\n"468 token_count += 1469 470 elapsed_ms = (time.perf_counter() - start_time) * 1000471 472 # Save bot response473 save_message(session_id, "bot", final_response)474 475 # Send done476 yield f"data: {json_module.dumps({'type': 'done', 'tokens': token_count, 'time_ms': elapsed_ms, 'session_id': session_id})}\n\n"477 478 return StreamingResponse(479 event_stream(),480 media_type="text/event-stream",481 headers={482 "Cache-Control": "no-cache, no-transform",483 "Connection": "keep-alive",484 "X-Accel-Buffering": "no",485 "X-Content-Type-Options": "nosniff"486 }487 )488 489 490# ── Legacy non-streaming (backward compat) ──491class ChatRequest(BaseModel):492 message: str493 session_id: Optional[str] = None494 history: Optional[list] = None495 temperature: float = 0.5496 top_k: int = 50497 top_p: float = 0.9498 repetition_penalty: float = 1.3499 max_tokens: int = 100500 501 502@app.post("/api/chat")503async def chat(req: ChatRequest):504 if model is None:505 raise HTTPException(status_code=503, detail="Model chưa được load.")506 507 message = req.message.strip()508 if not message:509 raise HTTPException(status_code=400, detail="Tin nhắn không được để trống.")510 511 response, token_count, elapsed_ms = generate_response(512 message, req.temperature, req.top_k, req.top_p,513 req.repetition_penalty, req.max_tokens514 )515 516 return {"response": response, "tokens_generated": token_count, "inference_time_ms": elapsed_ms}517 518 519# ── Static files ──520web_dir = os.path.join(os.path.dirname(__file__), "web")521if os.path.isdir(web_dir):522 app.mount("/assets", StaticFiles(directory=os.path.join(web_dir, "assets")), name="assets")523 524 @app.get("/")525 async def serve_frontend():526 return FileResponse(527 os.path.join(web_dir, "index.html"),528 headers={"Cache-Control": "no-cache, no-store, must-revalidate", "Pragma": "no-cache", "Expires": "0"}529 )530 531 532if __name__ == "__main__":533 import uvicorn534 uvicorn.run("api:app", host="127.0.0.1", port=8000, reload=True)535