Spaces:
Sleeping
Sleeping
File size: 4,204 Bytes
af92b4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """
Contextual Retrieval — sinh 1-2 câu context cho mỗi chunk dựa trên toàn văn tài liệu.
Context được prepend vào chunk trước khi embed + BM25, giúp retrieval chính xác hơn.
Tham khảo: Anthropic Contextual Retrieval paper (~35% giảm retrieval failure).
"""
import os
import google.generativeai as genai
from dotenv import load_dotenv
load_dotenv()
genai.configure(api_key=os.environ["GOOGLE_API_KEY"].strip())
_model = genai.GenerativeModel("gemini-2.0-flash")
CONTEXT_MODEL = "gemini-2.0-flash"
CONTEXT_PROMPT = """Dưới đây là toàn bộ tài liệu:
<document>
{doc_text}
</document>
Đây là đoạn văn cần đặt vào ngữ cảnh:
<chunk>
{chunk_text}
</chunk>
Hãy viết 1-2 câu ngắn gọn mô tả vị trí và vai trò của đoạn này trong tài liệu, \
giúp người đọc hiểu đoạn này nằm ở đâu và nói về điều gì.
Chỉ trả về 1-2 câu đó, không giải thích thêm."""
def generate_context(doc_text: str, chunk_text: str) -> str:
"""Sinh context cho 1 chunk. Trả về context string."""
# Giới hạn doc_text để tránh vượt context window
max_doc_chars = 8000
if len(doc_text) > max_doc_chars:
doc_text = doc_text[:max_doc_chars] + "\n...[nội dung tiếp theo]..."
import time
prompt = CONTEXT_PROMPT.format(doc_text=doc_text, chunk_text=chunk_text)
for attempt in range(5):
try:
response = _model.generate_content(
prompt,
generation_config={"max_output_tokens": 150, "temperature": 0.1},
)
return response.text.strip()
except Exception as e:
msg = str(e)
if "429" in msg:
# TPD (per day) exhausted — no point retrying
if "PerDay" in msg or "limit: 0" in msg:
print(f" [contextualizer] hết quota ngày, dừng sinh context")
raise RuntimeError("QUOTA_DAY_EXHAUSTED") from e
# TPM (per minute) — wait and retry
wait = 60 * (attempt + 1)
print(f" [contextualizer] rate limit, chờ {wait}s rồi thử lại...")
time.sleep(wait)
else:
print(f" [contextualizer] lỗi: {e}, bỏ qua context")
return ""
return ""
def add_context_to_chunks(chunks: list[dict], source_pages: list[dict]) -> list[dict]:
"""
Prepend context vào text của mỗi chunk.
source_pages: list các page dict có 'text' và 'metadata.source' để build doc_text.
"""
# Group pages theo source file
from collections import defaultdict
doc_texts = defaultdict(list)
for page in source_pages:
src = page["metadata"].get("source", "unknown")
doc_texts[src].append(page["text"])
# Full text per document
full_docs = {src: "\n\n".join(pages) for src, pages in doc_texts.items()}
import time
enriched = []
total = len(chunks)
for i, chunk in enumerate(chunks):
src = chunk["metadata"].get("source", "unknown")
doc_text = full_docs.get(src, "")
print(f" Generating context {i+1}/{total}: {src}...")
try:
context = generate_context(doc_text, chunk["text"])
except RuntimeError as e:
if "QUOTA_DAY_EXHAUSTED" in str(e):
print(f" [contextualizer] quota ngày hết tại chunk {i+1}/{total}, dừng lại")
# Append remaining chunks without context
for remaining in chunks[i:]:
r = dict(remaining)
r["metadata"] = {**remaining["metadata"], "has_context": False}
enriched.append(r)
return enriched
context = ""
time.sleep(4) # Gemini Flash free tier: 15 RPM
new_chunk = dict(chunk)
if context:
new_chunk["text"] = f"{context}\n\n{chunk['text']}"
new_chunk["metadata"] = {**chunk["metadata"], "has_context": True}
else:
new_chunk["metadata"] = {**chunk["metadata"], "has_context": False}
enriched.append(new_chunk)
return enriched
|