207 lines
7.0 KiB
Python
207 lines
7.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
SPIKE: atom extraction + bge-m3 embeddings + Qdrant dedup.
|
||
|
||
Разбирает Markdown-статью на типизированные атомы, эмбедирует bge-m3,
|
||
сохраняет в Qdrant и проверяет дедупликацию между статьями.
|
||
|
||
bge-m3 — единственная embed-модель (dim=1024). Не переключаться на nomic.
|
||
|
||
Запуск:
|
||
python3 spike/extract_atoms.py --article file.md [--article file2.md] --project rabota-ufa
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import time
|
||
import uuid
|
||
from dataclasses import dataclass, field
|
||
|
||
from qdrant_client import QdrantClient
|
||
from qdrant_client.models import Distance, VectorParams, PointStruct
|
||
|
||
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
|
||
EMBED_MODEL = "bge-m3"
|
||
COLLECTION = "content_atoms"
|
||
DIM = 1024
|
||
|
||
|
||
@dataclass
|
||
class Atom:
|
||
atom_id: str
|
||
project: str
|
||
type: str
|
||
title: str
|
||
summary: str
|
||
body: str
|
||
plain_text: str
|
||
word_count: int = 0
|
||
created_at: str = field(default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
|
||
|
||
def to_payload(self) -> dict:
|
||
return {
|
||
"atom_id": self.atom_id,
|
||
"project": self.project,
|
||
"type": self.type,
|
||
"title": self.title,
|
||
"summary": self.summary,
|
||
"body": self.body,
|
||
"plain_text": self.plain_text,
|
||
"word_count": self.word_count,
|
||
"created_at": self.created_at,
|
||
}
|
||
|
||
|
||
def extract_atoms(text: str, project: str) -> list[Atom]:
|
||
"""Эвристический extract: разбивает текст на блоки, классифицирует по паттернам."""
|
||
lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
|
||
atoms: list[Atom] = []
|
||
seq = 0
|
||
|
||
for ln in lines:
|
||
if ln.startswith("#") or ln.startswith(">") or ln.startswith("http"):
|
||
continue
|
||
if re.match(r"^\|.*\|$", ln) or ln.startswith("---") or ln.startswith("```"):
|
||
continue
|
||
if len(ln) < 20:
|
||
continue
|
||
seq += 1
|
||
atoms.append(Atom(
|
||
atom_id=f"{project}_{classify(ln)}_{seq:04d}",
|
||
project=project,
|
||
type=classify(ln),
|
||
title=ln[:80],
|
||
summary=ln[:160],
|
||
body=ln,
|
||
plain_text=ln,
|
||
word_count=len(re.findall(r"\S+", ln)),
|
||
))
|
||
return atoms
|
||
|
||
|
||
def classify(line: str) -> str:
|
||
low = line.lower()
|
||
if re.search(r"\d", line) and re.search(r"[%₽$€]", line):
|
||
return "stat"
|
||
if "?" in line:
|
||
return "qa"
|
||
if re.search(r"(подпишись|подписка|перейди|заказывай|оставь|свяжись|позвони)", low):
|
||
return "cta"
|
||
if re.search(r"(правильно|неправильно|ошибаеш|ты не|вот почему|секрет)", low):
|
||
return "hook"
|
||
if re.match(r"\d+[\.\)]", line):
|
||
return "instruction"
|
||
if line.startswith("«") or line.startswith("«"):
|
||
return "quote"
|
||
if re.search(r"(vs|против|по сравнению|лучше|дороже|дешевле)", low):
|
||
return "comparison"
|
||
return "claim"
|
||
|
||
|
||
def embed(text: str) -> list[float]:
|
||
"""bge-m3 embed через Ollama (dim=1024)."""
|
||
import urllib.request
|
||
body = json.dumps({"model": EMBED_MODEL, "prompt": text}).encode()
|
||
req = urllib.request.Request(
|
||
f"{OLLAMA_URL}/api/embeddings", data=body,
|
||
headers={"Content-Type": "application/json"},
|
||
)
|
||
with urllib.request.urlopen(req, timeout=60) as resp:
|
||
data = json.loads(resp.read().decode())
|
||
emb = data["embedding"]
|
||
if len(emb) != DIM:
|
||
raise ValueError(f"bge-m3 вернул dim={len(emb)}, ожидалось {DIM}")
|
||
return emb
|
||
|
||
|
||
def ensure_collection(client: QdrantClient) -> None:
|
||
if client.collection_exists(COLLECTION):
|
||
info = client.get_collection(COLLECTION)
|
||
print(f" ✓ коллекция {COLLECTION} уже существует (dim={info.config.params.vectors.size})")
|
||
return
|
||
|
||
client.create_collection(
|
||
collection_name=COLLECTION,
|
||
vectors_config=VectorParams(size=DIM, distance=Distance.COSINE),
|
||
)
|
||
print(f" ✓ коллекция {COLLECTION} создана (dim={DIM}, Cosine)")
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description="Spike: extract + bge-m3 + Qdrant dedup")
|
||
ap.add_argument("--article", action="append", help=".md файл (можно несколько)")
|
||
ap.add_argument("--project", default="spike", help="project_id (домен)")
|
||
ap.add_argument("--threshold", type=float, default=0.85, help="порог дедупа (Cosine)")
|
||
args = ap.parse_args()
|
||
|
||
if not args.article:
|
||
print("Использование: --article file.md [--article file2.md ...]")
|
||
return 1
|
||
|
||
client = QdrantClient(host="127.0.0.1", port=6333)
|
||
ensure_collection(client)
|
||
|
||
all_atoms: list[Atom] = []
|
||
for path in args.article:
|
||
if not os.path.exists(path):
|
||
print(f" ✗ файл не найден: {path}")
|
||
continue
|
||
with open(path, encoding="utf-8") as f:
|
||
text = f.read()
|
||
atoms = extract_atoms(text, args.project)
|
||
all_atoms.extend(atoms)
|
||
print(f" • {os.path.basename(path)}: {len(atoms)} атомов")
|
||
|
||
if not all_atoms:
|
||
print("Нет атомов для обработки")
|
||
return 1
|
||
|
||
print(f"\nЭмбеддинг bge-m3 ({len(all_atoms)} атомов)...")
|
||
start = time.time()
|
||
for i in range(0, len(all_atoms), 5):
|
||
for atom in all_atoms[i:i+5]:
|
||
emb = embed(atom.body)
|
||
client.upsert(
|
||
collection_name=COLLECTION,
|
||
points=[PointStruct(
|
||
id=str(uuid.uuid4()),
|
||
vector=emb,
|
||
payload=atom.to_payload(),
|
||
)],
|
||
)
|
||
print(f" • {min(i+5, len(all_atoms))}/{len(all_atoms)}")
|
||
elapsed = time.time() - start
|
||
print(f" ✓ загрузили {len(all_atoms)} атомов за {elapsed:.1f}с "
|
||
f"(~{elapsed/len(all_atoms)*1000:.0f}ms/atom)")
|
||
|
||
print(f"\nДедупликация (порог {args.threshold}):")
|
||
dedup_hits = 0
|
||
for atom in all_atoms:
|
||
emb = embed(atom.body)
|
||
res = client.query_points(
|
||
collection_name=COLLECTION,
|
||
query=emb,
|
||
limit=5,
|
||
with_payload=True,
|
||
)
|
||
for hit in res.points:
|
||
if hit.payload.get("atom_id") == atom.atom_id:
|
||
continue
|
||
score = hit.score or 0
|
||
if score >= args.threshold:
|
||
dedup_hits += 1
|
||
other = hit.payload.get("type", "?")
|
||
print(f" ⚠ ДУБЛЬ: {atom.type} ↔ {other} (score={score:.3f})")
|
||
break
|
||
|
||
print(f"\nИтог: {len(all_atoms)} атомов, {dedup_hits} дублей найдено")
|
||
print(f"Коллекция: 127.0.0.1:6333 / {COLLECTION}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|