vanhai123's picture
Upload 9 files
36cac86 verified
raw
history blame contribute delete
2.81 kB
import os
import requests
from retriever import retrieve_docs
from langchain_community.embeddings import HuggingFaceEmbeddings
from numpy import dot
from numpy.linalg import norm
API_KEY = "AIzaSyClqQssVMjt02qKrGKnghYAK9RkGf0lkS4"
def filter_relevant_docs(docs, query, top_k=3):
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
query_embedding = embeddings.embed_query(query)
scores = []
for doc in docs:
doc_embedding = embeddings.embed_query(doc.page_content)
cosine_sim = dot(query_embedding, doc_embedding) / (norm(query_embedding) * norm(doc_embedding))
scores.append((doc, cosine_sim))
scores.sort(key=lambda x: x[1], reverse=True)
return [doc for doc, _ in scores[:top_k]]
def format_sources(docs):
sources = set()
for doc in docs:
section = doc.metadata.get("section")
if section:
sources.add(section.strip())
else:
filename = os.path.basename(doc.metadata.get("source", "Nguồn không xác định"))
sources.add(filename)
return "\n".join(f"- {src}" for src in sorted(sources))
def answer_query(query, model="Gemini Pro", temperature=0.2):
all_docs = retrieve_docs(query)
if not all_docs:
return "Không tìm thấy tài liệu liên quan để trả lời.", []
docs = filter_relevant_docs(all_docs, query)
context = "\n\n".join([doc.page_content for doc in docs])
prompt = f"""Dựa trên tài liệu sau, hãy trả lời câu hỏi theo phong cách trang trọng, lịch sự và chuyên nghiệp:
{context}
Câu hỏi: {query}
Yêu cầu:
- Sử dụng từ ngữ lịch sự ("Bạn cần...", "Vui lòng...", "Sau khi...")
- Không sử dụng từ nói miệng như "nhé", "nha", "ok".
- Câu trúc câu đầy đủ, rõ ràng.
- Chỉ trả lời dựa trên thông tin trong tài liệu. Nếu không tìm thấy thông tin liên quan, trả lời: "Thông tin không có trong tài liệu được cung cấp."
- Không tự thêm "Nguồn tham khảo" trong phần trả lời.
Trả lời:"""
url = f"https://generativelanguage.googleapis.com/v1/models/gemini-1.5-pro:generateContent?key={API_KEY}"
headers = {"Content-Type": "application/json"}
payload = {
"contents": [{"parts": [{"text": prompt}]}],
"generationConfig": {"temperature": temperature}
}
response = requests.post(url, headers=headers, json=payload)
data = response.json()
try:
answer = data['candidates'][0]['content']['parts'][0]['text']
except Exception as e:
print("🔴 Response từ Gemini:", data)
answer = "Lỗi khi gọi Gemini API: " + str(e)
return answer, docs