chatbot_demo / rag_pipeline.py
deddoggo's picture
update
0c16fc9
raw
history blame
2.04 kB
# llm_handler.py
# Chịu trách nhiệm cho mọi logic liên quan đến mô hình ngôn ngữ lớn (LLM).
import torch
from unsloth import FastLanguageModel
# --- HÀM TẠO CÂU TRẢ LỜI ---
def generate_llm_response(
query: str,
context: str,
llm_model,
tokenizer,
max_new_tokens: int = 512,
temperature: float = 0.3,
top_p: float = 0.9,
) -> str:
"""
Sinh câu trả lời từ LLM dựa trên câu hỏi và ngữ cảnh đã được truy xuất.
"""
print("🧠 Bắt đầu sinh câu trả lời từ LLM...")
# Xây dựng prompt
prompt = f"""Bạn là một trợ lý AI chuyên tư vấn về luật giao thông đường bộ Việt Nam.
Dựa vào các thông tin luật được cung cấp dưới đây để trả lời câu hỏi của người dùng một cách chính xác và chi tiết.
Nếu thông tin không đủ, hãy trả lời rằng bạn không tìm thấy thông tin cụ thể trong tài liệu.
### Thông tin luật được trích dẫn:
{context}
### Câu hỏi của người dùng:
{query}
### Trả lời của bạn:"""
# Tạo input cho model
device = llm_model.device
inputs = tokenizer(prompt, return_tensors="pt").to(device)
# Cấu hình cho việc sinh văn bản
generation_config = dict(
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
try:
output_ids = llm_model.generate(**inputs, **generation_config)
input_length = inputs.input_ids.shape[1]
generated_ids = output_ids[0][input_length:]
response_text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
print("✅ Sinh câu trả lời hoàn tất.")
return response_text
except Exception as e:
print(f"❌ Lỗi khi sinh câu trả lời từ LLM: {e}")
return "Xin lỗi, đã có lỗi xảy ra trong quá trình tạo câu trả lời."