283 lines
9.4 KiB
Python
283 lines
9.4 KiB
Python
"""
|
|
Celery tasks for downloading and indexing document sources.
|
|
|
|
Redis key "index_job" stores current job state as JSON:
|
|
{ status, source, progress, message, started_at, finished_at, docs_total, docs_done }
|
|
"""
|
|
import bz2
|
|
import json
|
|
import os
|
|
import time
|
|
import requests
|
|
import xml.etree.ElementTree as ET
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
import redis
|
|
from app.celery_app import celery_app
|
|
|
|
from app.config import settings
|
|
from app.models.db import SessionLocal, Document, Fingerprint
|
|
|
|
JOB_KEY = "index_job"
|
|
|
|
WIKIPEDIA_SOURCES = {
|
|
"wikipedia_ru": {
|
|
"url": "https://dumps.wikimedia.org/ruwiki/latest/ruwiki-latest-pages-articles.xml.bz2",
|
|
"meta_url": "https://dumps.wikimedia.org/ruwiki/latest/ruwiki-latest-md5sums.txt",
|
|
"lang": "ru",
|
|
"label": "Wikipedia RU",
|
|
},
|
|
"wikipedia_en": {
|
|
"url": "https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2",
|
|
"meta_url": "https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-md5sums.txt",
|
|
"lang": "en",
|
|
"label": "Wikipedia EN",
|
|
},
|
|
}
|
|
|
|
|
|
# ── Redis helpers ──────────────────────────────────────────────────────────
|
|
|
|
def _r() -> redis.Redis:
|
|
return redis.from_url(settings.REDIS_URL)
|
|
|
|
|
|
def _set_status(data: dict):
|
|
_r().set(JOB_KEY, json.dumps(data, default=str))
|
|
|
|
|
|
def _get_status() -> dict:
|
|
raw = _r().get(JOB_KEY)
|
|
if not raw:
|
|
return {"status": "idle"}
|
|
return json.loads(raw)
|
|
|
|
|
|
def _pub(message: str, progress: int = 0, **extra):
|
|
state = _get_status()
|
|
state.update({"message": message, "progress": progress, **extra})
|
|
_set_status(state)
|
|
_r().publish("index_job_progress", json.dumps(state, default=str))
|
|
|
|
|
|
# ── Text helpers ───────────────────────────────────────────────────────────
|
|
|
|
import re
|
|
|
|
|
|
def _clean_wiki(text: str) -> str:
|
|
text = re.sub(r'\[\[([^\|]*\|)?([^\]]*)\]\]', r'\2', text)
|
|
text = re.sub(r'\[http[^\s]+ ([^\]]*)\]', r'\1', text)
|
|
text = re.sub(r'\{\{[^}]*\}\}', '', text)
|
|
text = re.sub(r'\[\[Category:[^\]]*\]\]', '', text)
|
|
text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL)
|
|
text = re.sub(r'<[^>]+>', '', text)
|
|
text = re.sub(r'\n{3,}', '\n\n', text)
|
|
text = re.sub(r' +', ' ', text)
|
|
return text.strip()
|
|
|
|
|
|
def _normalize(text: str) -> str:
|
|
text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f]', '', text)
|
|
text = re.sub(r'[—–‒―]', '-', text)
|
|
text = re.sub(r'[ \t]+', ' ', text)
|
|
text = re.sub(r'\n{3,}', '\n\n', text)
|
|
return text.strip()
|
|
|
|
|
|
# ── Download helper ────────────────────────────────────────────────────────
|
|
|
|
def _download_file(url: str, dest: Path, label: str) -> Path:
|
|
"""Stream-download url → dest, publishing progress."""
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
r = requests.get(url, stream=True, timeout=30)
|
|
r.raise_for_status()
|
|
total = int(r.headers.get("Content-Length", 0))
|
|
done = 0
|
|
last_pct = -1
|
|
|
|
with open(dest, "wb") as f:
|
|
for chunk in r.iter_content(chunk_size=1 << 20): # 1 MB
|
|
f.write(chunk)
|
|
done += len(chunk)
|
|
pct = int(done / total * 40) if total else 0 # 0-40% of overall
|
|
if pct != last_pct:
|
|
_pub(f"Скачивание {label}: {done >> 20} MB", progress=pct)
|
|
last_pct = pct
|
|
|
|
return dest
|
|
|
|
|
|
# ── Indexing helper ────────────────────────────────────────────────────────
|
|
|
|
def _index_wikipedia_bz2(bz2_path: Path, source: str, lang: str, update_only: bool) -> int:
|
|
"""
|
|
Parse Wikipedia bz2 dump and index documents.
|
|
update_only=True → skip ext_ids already in DB.
|
|
Returns count of newly indexed docs.
|
|
"""
|
|
from app.core.fingerprint import WinNowing, MinHashLSHIndex
|
|
from app.core.embeddings import EmbeddingIndex
|
|
from app.core.fulltext import FulltextSearch
|
|
|
|
db = SessionLocal()
|
|
winnowing = WinNowing()
|
|
minhash = MinHashLSHIndex()
|
|
emb = EmbeddingIndex()
|
|
fulltext = FulltextSearch()
|
|
|
|
# Pre-load known ext_ids for fast skip
|
|
known: set[str] = set()
|
|
if update_only:
|
|
for (ext_id,) in db.query(Document.ext_id).filter(Document.source == source).yield_per(10_000):
|
|
if ext_id:
|
|
known.add(ext_id)
|
|
|
|
ns = {'wiki': 'http://www.mediawiki.org/xml/export-0.10/'}
|
|
|
|
embed_batch: list[tuple[str, str]] = []
|
|
BATCH = 32
|
|
count = 0
|
|
parsed = 0
|
|
|
|
with bz2.BZ2File(str(bz2_path), 'r') as f:
|
|
context = ET.iterparse(f, events=('end',))
|
|
for event, elem in context:
|
|
if not elem.tag.endswith('page'):
|
|
continue
|
|
|
|
title_el = elem.find('wiki:title', ns) or elem.find('.//title')
|
|
id_el = elem.find('wiki:id', ns) or elem.find('.//id')
|
|
rev = elem.find('.//revision')
|
|
text_el = rev.find('wiki:text', ns) if rev is not None else None
|
|
if rev is None:
|
|
text_el = elem.find('.//text')
|
|
|
|
if title_el is None or text_el is None:
|
|
elem.clear(); continue
|
|
|
|
title = (title_el.text or '').strip()
|
|
page_id = (id_el.text or '').strip() if id_el is not None else ''
|
|
raw_text = text_el.text or ''
|
|
ext_id = f"wiki_{page_id}"
|
|
|
|
if raw_text.startswith('#REDIRECT') or title.startswith('Wikipedia:'):
|
|
elem.clear(); continue
|
|
|
|
text = _normalize(_clean_wiki(raw_text))
|
|
if len(text.split()) < 50:
|
|
elem.clear(); continue
|
|
|
|
parsed += 1
|
|
|
|
if update_only and ext_id in known:
|
|
elem.clear(); continue
|
|
|
|
url = f"https://{lang}.wikipedia.org/wiki/{title.replace(' ', '_')}"
|
|
|
|
# PostgreSQL
|
|
doc = Document(source=source, ext_id=ext_id, title=title, url=url, lang=lang)
|
|
db.add(doc)
|
|
db.flush()
|
|
|
|
# Fingerprints
|
|
fp = winnowing.fingerprint(text)
|
|
for pos, h in enumerate(list(fp)[:500]):
|
|
db.add(Fingerprint(doc_id=doc.id, hash=format(h, '016x'), position=pos))
|
|
|
|
# MinHash
|
|
try:
|
|
minhash.add_document(str(doc.id), text)
|
|
except ValueError:
|
|
pass
|
|
|
|
# Elasticsearch
|
|
fulltext.index_document(str(doc.id), title, text, source, lang, url)
|
|
|
|
embed_batch.append((str(doc.id), text))
|
|
if len(embed_batch) >= BATCH:
|
|
emb.batch_add(embed_batch, batch_size=BATCH)
|
|
embed_batch.clear()
|
|
|
|
count += 1
|
|
if count % 500 == 0:
|
|
db.commit()
|
|
pct = 40 + min(int(count / 100000 * 55), 55) # 40-95%
|
|
_pub(f"Проиндексировано: {count} документов", progress=pct, docs_done=count)
|
|
|
|
elem.clear()
|
|
|
|
if embed_batch:
|
|
emb.batch_add(embed_batch, batch_size=BATCH)
|
|
|
|
db.commit()
|
|
db.close()
|
|
emb.save_index()
|
|
|
|
return count
|
|
|
|
|
|
# ── Celery tasks ───────────────────────────────────────────────────────────
|
|
|
|
def _run_indexing(source_key: str, update_only: bool):
|
|
src = WIKIPEDIA_SOURCES[source_key]
|
|
label = src["label"]
|
|
lang = src["lang"]
|
|
|
|
_set_status({
|
|
"status": "running",
|
|
"source": source_key,
|
|
"label": label,
|
|
"progress": 0,
|
|
"message": "Запуск…",
|
|
"started_at": datetime.now(timezone.utc).isoformat(),
|
|
"docs_done": 0,
|
|
})
|
|
|
|
data_dir = Path(settings.RAW_DATA_DIR) / source_key
|
|
bz2_path = data_dir / "dump.xml.bz2"
|
|
|
|
# --- Download if needed ---
|
|
if not bz2_path.exists() or not update_only:
|
|
_pub(f"Скачивание дампа {label}…", progress=1)
|
|
try:
|
|
_download_file(src["url"], bz2_path, label)
|
|
except Exception as e:
|
|
_set_status({"status": "error", "message": f"Ошибка скачивания: {e}"})
|
|
return
|
|
|
|
_pub("Индексация…", progress=41)
|
|
|
|
try:
|
|
count = _index_wikipedia_bz2(bz2_path, source_key, lang, update_only)
|
|
except Exception as e:
|
|
_set_status({"status": "error", "message": f"Ошибка индексации: {e}"})
|
|
return
|
|
|
|
action = "обновлено" if update_only else "проиндексировано"
|
|
_set_status({
|
|
"status": "done",
|
|
"source": source_key,
|
|
"label": label,
|
|
"progress": 100,
|
|
"message": f"Готово — {action} {count} документов",
|
|
"finished_at": datetime.now(timezone.utc).isoformat(),
|
|
"docs_done": count,
|
|
})
|
|
_r().publish("index_job_progress", _r().get(JOB_KEY))
|
|
|
|
|
|
@celery_app.task(name="tasks.download_and_index", bind=True)
|
|
def download_and_index(self, source_key: str):
|
|
_run_indexing(source_key, update_only=False)
|
|
|
|
|
|
@celery_app.task(name="tasks.update_index", bind=True)
|
|
def update_index(self, source_key: str):
|
|
_run_indexing(source_key, update_only=True)
|
|
|
|
|
|
def get_job_status() -> dict:
|
|
return _get_status()
|