Upload chat.py with huggingface_hub

3051a6d verified 5 months ago

7.48 kB

"""
Chat trực tiếp qua terminal.
Tích hợp BPE Tokenizer + Cosine Similarity.
Chạy: python chat.py
"""
import torch
import torch.nn.functional as F
import pickle
import os
import re
import sys
import csv
from src.model import make_model

sys.stdout.reconfigure(encoding='utf-8')

VOCAB_PATH     = 'data/processed/vocab.pkl'
MODEL_PATH     = 'models/best/model.pt'
TOKENIZER_PATH = 'data/processed/bpe_tokenizer.json'
RAW_CSV_PATH   = 'data/raw/output.csv'

SIMILARITY_THRESHOLD = 0.92


def load_bpe_tokenizer():
    if os.path.exists(TOKENIZER_PATH):
        from tokenizers import Tokenizer
        return Tokenizer.from_file(TOKENIZER_PATH)
    return None


def tokenize_bpe(text, bpe_tok, word2idx):
    if bpe_tok is not None:
        return bpe_tok.encode(str(text).strip()).ids
    else:
        text = str(text).lower()
        text = re.sub(r'([.,!?()])', r' \1 ', text)
        text = re.sub(r'\s{2,}', ' ', text)
        return [word2idx.get(w, word2idx.get('<unk>', 3)) for w in text.strip().split()]


def build_qa_embeddings(model, device, bpe_tok, word2idx):
    qa_database = []
    if not os.path.exists(RAW_CSV_PATH):
        return qa_database, None

    with open(RAW_CSV_PATH, 'r', encoding='utf-8-sig') as f:
        reader = csv.reader(f)
        next(reader, None)
        for row in reader:
            if len(row) >= 3 and row[1].strip() and row[2].strip():
                qa_database.append({
                    "topic": row[0].strip(),
                    "question": row[1].strip(),
                    "answer": row[2].strip(),
                })

    if not qa_database:
        return qa_database, None

    sos_id = word2idx.get('<sos>', 1)
    all_embs = []
    for item in qa_database:
        q_ids = tokenize_bpe(item['question'], bpe_tok, word2idx)
        x = torch.tensor([[sos_id] + q_ids]).to(device)
        emb = model.get_embedding(x)
        all_embs.append(emb)

    qa_embeddings = F.normalize(torch.cat(all_embs, dim=0), p=2, dim=1)
    return qa_database, qa_embeddings


def find_similar(user_q, model, device, bpe_tok, word2idx, qa_database, qa_embeddings):
    if qa_embeddings is None:
        return None, 0.0
    sos_id = word2idx.get('<sos>', 1)
    q_ids = tokenize_bpe(user_q, bpe_tok, word2idx)
    x = torch.tensor([[sos_id] + q_ids]).to(device)
    user_emb = F.normalize(model.get_embedding(x), p=2, dim=1)
    sims = (qa_embeddings @ user_emb.T).squeeze(1)
    best_idx = sims.argmax().item()
    best_score = sims[best_idx].item()
    if best_score >= SIMILARITY_THRESHOLD:
        return qa_database[best_idx], best_score
    return None, best_score


def clean_response(text):
    text = text.strip()
    if not text:
        return ""
    # Viết hoa chữ cái đầu tiên
    text = text[0].upper() + text[1:]
    
    # Chuẩn hóa Unicode (sửa lỗi dấu tiếng Việt NFD nếu có)
    import unicodedata
    text = unicodedata.normalize('NFC', text)
    
    # Sửa lỗi khoảng trắng trong LaTeX (do Tokenizer tách)
    text = re.sub(r'\\\s+([a-zA-Z]+)', r'\\\1', text) # \ cdot -> \cdot
    text = re.sub(r'\s*_\s*', '_', text)            # k _ 5 -> k_5
    text = re.sub(r'\s*\^\s*', '^', text)            # x ^ 2 -> x^2
    text = re.sub(r'\s*\{\s*', '{', text)            # { 5 } -> {5}
    text = re.sub(r'\s*\}\s*', '}', text)
    
    # Xóa khoảng trắng thừa trước các dấu câu
    text = re.sub(r'\s+([.,!?():;])', r'\1', text)
    # Xóa khoảng trắng thừa sau dấu mở ngoặc
    text = re.sub(r'\(\s+', r'(', text)
    # Rút gọn khoảng trắng trùng lặp
    text = re.sub(r'\s{2,}', ' ', text)
    return text.strip()


def chat():
    if not os.path.exists(MODEL_PATH):
        print(f"Không tìm thấy model: {MODEL_PATH}")
        print("Hãy chạy train_gpu.py trước.")
        return

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    with open(VOCAB_PATH, 'rb') as f:
        vocab = pickle.load(f)
    word2idx = vocab['word2idx']
    idx2word = vocab['idx2word']

    ckpt  = torch.load(MODEL_PATH, map_location=device)
    
    # Hỗ trợ load cả model.pt và checkpoint_large.pt
    d_m = ckpt.get('d_m', 768)
    n_l = ckpt.get('n_l', 12)
    n_h = ckpt.get('n_h', 12)
    dr  = ckpt.get('dr', 0.1)

    model = make_model(
        len(word2idx),
        n_layer=n_l, d_model=d_m,
        n_head=n_h,  dropout=dr
    )
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device)
    model.eval()

    bpe_tok = load_bpe_tokenizer()
    print(f"  BPE Tokenizer: {'Loaded' if bpe_tok else 'Fallback to split()'}")

    # Đã tắt để AI tự sinh 100% nhằm tăng tốc khởi động
    qa_database, qa_embeddings = [], None
    print("  [INFO] Bỏ qua build embeddings vì đã tắt Similarity Match.")

    print(f"\n--- CHATBOT TOÁN RỜI RẠC | device={device} ---")
    print("Gõ 'exit' để thoát.\n")

    while True:
        try:
            user_input = input("Bạn: ").strip()
            if user_input.lower() in ['exit', 'quit', 'thoát']:
                break
            if not user_input:
                continue

            # Bước 1: So khớp tương đồng (Đã tắt để AI tự sinh 100%)
            matched, score = None, 0.0
            if matched is not None:
                print(f"Bot: {matched['answer']}")
                print(f"  [Similarity: {score:.2%} | Matched: \"{matched['question']}\"]")
                print()
                continue

            # Bước 2: Generate
            tokens = tokenize_bpe(user_input, bpe_tok, word2idx)
            sos_id = word2idx.get('<sos>', 1)
            sep_id = word2idx.get('<sep>', 4)
            prompt = [sos_id] + tokens + [sep_id]
            x = torch.tensor([prompt]).to(device)

            with torch.no_grad():
                out_ids = model.generate(x, max_new_tokens=100,
                                         temperature=0.5, top_k=50, top_p=0.9,
                                         repetition_penalty=1.3)

            resp_ids = out_ids[0][len(prompt):]

            if bpe_tok is not None:
                valid_ids = []
                for i in resp_ids:
                    tid = i.item()
                    if tid == word2idx.get('<eos>', 2):
                        break
                    if tid not in [word2idx.get('<pad>', 0), word2idx.get('<sos>', 1),
                                   word2idx.get('<unk>', 3), word2idx.get('<sep>', 4)]:
                        valid_ids.append(tid)
                response = bpe_tok.decode(valid_ids)
            else:
                resp_tokens = []
                for i in resp_ids:
                    sym = idx2word.get(i.item(), '<unk>')
                    if sym == '<eos>':
                        break
                    if sym not in ['<pad>', '<sos>', '<unk>', '<sep>']:
                        resp_tokens.append(sym)
                response = " ".join(resp_tokens)

            # Làm sạch câu trả lời sinh ra
            response = clean_response(response)

            if not response.strip():
                response = "Mình chưa học tới phần này. Bạn hỏi câu khác nhé!"

            print(f"Bot: {response}")
            print(f"  [Generated | Similarity best: {score:.2%}]")
            print()

        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"Lỗi: {e}")


if __name__ == "__main__":
    chat()

	"""
	Chat trực tiếp qua terminal.
	Tích hợp BPE Tokenizer + Cosine Similarity.
	Chạy: python chat.py
	"""
	import torch
	import torch.nn.functional as F
	import pickle
	import os
	import re
	import sys
	import csv
	from src.model import make_model

	sys.stdout.reconfigure(encoding='utf-8')

	VOCAB_PATH = 'data/processed/vocab.pkl'
	MODEL_PATH = 'models/best/model.pt'
	TOKENIZER_PATH = 'data/processed/bpe_tokenizer.json'
	RAW_CSV_PATH = 'data/raw/output.csv'

	SIMILARITY_THRESHOLD = 0.92


	def load_bpe_tokenizer():
	if os.path.exists(TOKENIZER_PATH):
	from tokenizers import Tokenizer
	return Tokenizer.from_file(TOKENIZER_PATH)
	return None


	def tokenize_bpe(text, bpe_tok, word2idx):
	if bpe_tok is not None:
	return bpe_tok.encode(str(text).strip()).ids
	else:
	text = str(text).lower()
	text = re.sub(r'([.,!?()])', r' \1 ', text)
	text = re.sub(r'\s{2,}', ' ', text)
	return [word2idx.get(w, word2idx.get('<unk>', 3)) for w in text.strip().split()]


	def build_qa_embeddings(model, device, bpe_tok, word2idx):
	qa_database = []
	if not os.path.exists(RAW_CSV_PATH):
	return qa_database, None

	with open(RAW_CSV_PATH, 'r', encoding='utf-8-sig') as f:
	reader = csv.reader(f)
	next(reader, None)
	for row in reader:
	if len(row) >= 3 and row[1].strip() and row[2].strip():
	qa_database.append({
	"topic": row[0].strip(),
	"question": row[1].strip(),
	"answer": row[2].strip(),
	})

	if not qa_database:
	return qa_database, None

	sos_id = word2idx.get('<sos>', 1)
	all_embs = []
	for item in qa_database:
	q_ids = tokenize_bpe(item['question'], bpe_tok, word2idx)
	x = torch.tensor([[sos_id] + q_ids]).to(device)
	emb = model.get_embedding(x)
	all_embs.append(emb)

	qa_embeddings = F.normalize(torch.cat(all_embs, dim=0), p=2, dim=1)
	return qa_database, qa_embeddings


	def find_similar(user_q, model, device, bpe_tok, word2idx, qa_database, qa_embeddings):
	if qa_embeddings is None:
	return None, 0.0
	sos_id = word2idx.get('<sos>', 1)
	q_ids = tokenize_bpe(user_q, bpe_tok, word2idx)
	x = torch.tensor([[sos_id] + q_ids]).to(device)
	user_emb = F.normalize(model.get_embedding(x), p=2, dim=1)
	sims = (qa_embeddings @ user_emb.T).squeeze(1)
	best_idx = sims.argmax().item()
	best_score = sims[best_idx].item()
	if best_score >= SIMILARITY_THRESHOLD:
	return qa_database[best_idx], best_score
	return None, best_score


	def clean_response(text):
	text = text.strip()
	if not text:
	return ""
	# Viết hoa chữ cái đầu tiên
	text = text[0].upper() + text[1:]

	# Chuẩn hóa Unicode (sửa lỗi dấu tiếng Việt NFD nếu có)
	import unicodedata
	text = unicodedata.normalize('NFC', text)

	# Sửa lỗi khoảng trắng trong LaTeX (do Tokenizer tách)
	text = re.sub(r'\\\s+([a-zA-Z]+)', r'\\\1', text) # \ cdot -> \cdot
	text = re.sub(r'\s_\s', '_', text) # k _ 5 -> k_5
	text = re.sub(r'\s\^\s', '^', text) # x ^ 2 -> x^2
	text = re.sub(r'\s\{\s', '{', text) # { 5 } -> {5}
	text = re.sub(r'\s\}\s', '}', text)

	# Xóa khoảng trắng thừa trước các dấu câu
	text = re.sub(r'\s+([.,!?():;])', r'\1', text)
	# Xóa khoảng trắng thừa sau dấu mở ngoặc
	text = re.sub(r'\(\s+', r'(', text)
	# Rút gọn khoảng trắng trùng lặp
	text = re.sub(r'\s{2,}', ' ', text)
	return text.strip()


	def chat():
	if not os.path.exists(MODEL_PATH):
	print(f"Không tìm thấy model: {MODEL_PATH}")
	print("Hãy chạy train_gpu.py trước.")
	return

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	with open(VOCAB_PATH, 'rb') as f:
	vocab = pickle.load(f)
	word2idx = vocab['word2idx']
	idx2word = vocab['idx2word']

	ckpt = torch.load(MODEL_PATH, map_location=device)

	# Hỗ trợ load cả model.pt và checkpoint_large.pt
	d_m = ckpt.get('d_m', 768)
	n_l = ckpt.get('n_l', 12)
	n_h = ckpt.get('n_h', 12)
	dr = ckpt.get('dr', 0.1)

	model = make_model(
	len(word2idx),
	n_layer=n_l, d_model=d_m,
	n_head=n_h, dropout=dr
	)
	model.load_state_dict(ckpt['model_state_dict'])
	model.to(device)
	model.eval()

	bpe_tok = load_bpe_tokenizer()
	print(f" BPE Tokenizer: {'Loaded' if bpe_tok else 'Fallback to split()'}")

	# Đã tắt để AI tự sinh 100% nhằm tăng tốc khởi động
	qa_database, qa_embeddings = [], None
	print(" [INFO] Bỏ qua build embeddings vì đã tắt Similarity Match.")

	print(f"\n--- CHATBOT TOÁN RỜI RẠC \| device={device} ---")
	print("Gõ 'exit' để thoát.\n")

	while True:
	try:
	user_input = input("Bạn: ").strip()
	if user_input.lower() in ['exit', 'quit', 'thoát']:
	break
	if not user_input:
	continue

	# Bước 1: So khớp tương đồng (Đã tắt để AI tự sinh 100%)
	matched, score = None, 0.0
	if matched is not None:
	print(f"Bot: {matched['answer']}")
	print(f" [Similarity: {score:.2%} \| Matched: \"{matched['question']}\"]")
	print()
	continue

	# Bước 2: Generate
	tokens = tokenize_bpe(user_input, bpe_tok, word2idx)
	sos_id = word2idx.get('<sos>', 1)
	sep_id = word2idx.get('<sep>', 4)
	prompt = [sos_id] + tokens + [sep_id]
	x = torch.tensor([prompt]).to(device)

	with torch.no_grad():
	out_ids = model.generate(x, max_new_tokens=100,
	temperature=0.5, top_k=50, top_p=0.9,
	repetition_penalty=1.3)

	resp_ids = out_ids[0][len(prompt):]

	if bpe_tok is not None:
	valid_ids = []
	for i in resp_ids:
	tid = i.item()
	if tid == word2idx.get('<eos>', 2):
	break
	if tid not in [word2idx.get('<pad>', 0), word2idx.get('<sos>', 1),
	word2idx.get('<unk>', 3), word2idx.get('<sep>', 4)]:
	valid_ids.append(tid)
	response = bpe_tok.decode(valid_ids)
	else:
	resp_tokens = []
	for i in resp_ids:
	sym = idx2word.get(i.item(), '<unk>')
	if sym == '<eos>':
	break
	if sym not in ['<pad>', '<sos>', '<unk>', '<sep>']:
	resp_tokens.append(sym)
	response = " ".join(resp_tokens)

	# Làm sạch câu trả lời sinh ra
	response = clean_response(response)

	if not response.strip():
	response = "Mình chưa học tới phần này. Bạn hỏi câu khác nhé!"

	print(f"Bot: {response}")
	print(f" [Generated \| Similarity best: {score:.2%}]")
	print()

	except KeyboardInterrupt:
	break
	except Exception as e:
	print(f"Lỗi: {e}")


	if __name__ == "__main__":
	chat()