turkish-wikipedia-rag / data_preparation.py
ecceembusra's picture
Upload 5 files
6dda1eb verified
raw
history blame
5.35 kB
# data_preparation.py
import os, re, time, json
from typing import List, Dict, Tuple
import numpy as np
import faiss
from datasets import load_dataset, DownloadConfig
from langchain_text_splitters import RecursiveCharacterTextSplitter
# Embedding sağlayıcısı (E5) -> providers.embed
from providers import embed
# =========================
# AYARLAR
# =========================
DATASET_NAME = "wikimedia/wikipedia" # HF dataset adı
WIKI_CONFIG = "20231101.tr" # Türkçe dump sürümü
MAX_PAGES = 500 # hızlı demo için 500
CHUNK_SIZE = 1000 # büyük parça = daha az chunk
CHUNK_OVERLAP= 100
MAX_CHUNKS = 8000 # güvenli üst limit (None yapıldığında sınırsız)
VSTORE_DIR = "vectorstore"
META_JSONL = "meta.jsonl"
FAISS_FILE = "index.faiss"
SIGN_FILE = "signature.json"
# =========================
# Yardımcılar
# =========================
def ensure_dir(p: str):
os.makedirs(p, exist_ok=True)
def slugify_title(title: str) -> str:
t = (title or "").strip().replace(" ", "_")
t = re.sub(r"[^\w\-ÇçĞğİıÖöŞşÜü]+", "", t, flags=re.UNICODE)
return t or "Sayfa"
# Gereksiz template token temizliği
TOKEN_PAT = re.compile(r"<\|/?(system|user|assistant|start|end)\|>|<\/?s>|<\/?unk>", re.IGNORECASE)
def clean_text(s: str) -> str:
if not s:
return ""
s = TOKEN_PAT.sub(" ", s)
return re.sub(r"\s+", " ", s).strip()
# =========================
# Chunking
# =========================
def chunk_documents(rows: List[Dict]) -> Tuple[List[str], List[Dict]]:
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
)
texts_raw: List[str] = []
metas: List[Dict] = []
for r in rows:
title = clean_text(r.get("title", ""))
text = clean_text(r.get("text", ""))
if not text:
continue
url = f"https://tr.wikipedia.org/wiki/{slugify_title(title)}"
chunks = splitter.split_text(text)
for i, ch in enumerate(chunks):
ch = ch.strip()
if not ch:
continue
texts_raw.append(ch) # meta'ya RAW metin yazılıyor
metas.append({"title": title or "(başlık yok)", "chunk_id": i, "source": url})
if MAX_CHUNKS and len(texts_raw) >= MAX_CHUNKS:
return texts_raw, metas
return texts_raw, metas
# =========================
# FAISS (HNSW) — hızlı ANN arama
# =========================
def build_faiss_index(vecs: np.ndarray) -> faiss.Index:
"""
Cosine benzerliği için L2 normalize edip Inner-Product ile HNSW kullanılır.
HNSW, IndexFlat'e yakın doğrulukta olup sorgu süresini ciddi düşürür.
"""
faiss.normalize_L2(vecs)
dim = vecs.shape[1]
M = 32 # graph degree
index = faiss.IndexHNSWFlat(dim, M, faiss.METRIC_INNER_PRODUCT)
index.hnsw.efConstruction = 80 # inşa kalitesi (yüksek = daha iyi/az fark)
index.add(vecs)
return index
# =========================
# Ana akış
# =========================
def main():
t0 = time.time()
print("👉 Wikipedia(TR) yükleniyor...")
split_expr = f"train[:{MAX_PAGES}]" if MAX_PAGES else "train"
ds = load_dataset(
DATASET_NAME, WIKI_CONFIG,
split=split_expr,
download_config=DownloadConfig(max_retries=5),
)
print(f"Toplam sayfa (seçim sonrası): {len(ds)}")
print("👉 Chunk'lanıyor...")
texts_raw, metas = chunk_documents([dict(x) for x in ds])
print(f"Toplam chunk: {len(texts_raw)}")
if not texts_raw:
raise SystemExit("⚠️ Metin bulunamadı.")
print("👉 Embedding hesaplanıyor (E5)...")
# E5 kuralı: embed edilirken PASSAGE prefix kullan, meta'da görünmesin
texts_for_emb = [f"passage: {t}" for t in texts_raw]
vecs = np.asarray(embed(texts_for_emb), dtype="float32")
if vecs.ndim != 2:
raise ValueError(f"Beklenen (N,D) vektör, gelen {vecs.shape}")
print("👉 FAISS (HNSW) indeks oluşturuluyor...")
index = build_faiss_index(vecs)
print("👉 Kaydediliyor...")
ensure_dir(VSTORE_DIR)
faiss.write_index(index, os.path.join(VSTORE_DIR, FAISS_FILE))
meta_path = os.path.join(VSTORE_DIR, META_JSONL)
with open(meta_path, "w", encoding="utf-8") as f:
for t, m in zip(texts_raw, metas):
f.write(json.dumps({"text": t, "metadata": m}, ensure_ascii=False) + "\n")
sign_path = os.path.join(VSTORE_DIR, SIGN_FILE)
with open(sign_path, "w", encoding="utf-8") as f:
json.dump({
"dataset": f"{DATASET_NAME}:{WIKI_CONFIG}",
"max_pages": MAX_PAGES,
"chunk_size": CHUNK_SIZE,
"chunk_overlap": CHUNK_OVERLAP,
"max_chunks": MAX_CHUNKS,
"faiss": {"type": "HNSWFlat", "metric": "IP", "M": 32, "efConstruction": 80},
"emb_model": os.getenv("EMB_MODEL", "intfloat/multilingual-e5-small"),
}, f, ensure_ascii=False, indent=2)
print(f"✅ Tamamlandı. Süre: {(time.time()-t0):.1f} sn | Çıktı klasörü: {VSTORE_DIR}")
if __name__ == "__main__":
main()