File size: 3,536 Bytes
325b94c | 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 112 113 114 115 116 117 118 119 120 121 122 | import os
import uuid
from dataclasses import asdict
from typing import List, Optional
from sentence_transformers import SentenceTransformer
from qdrant_client import QdrantClient
from qdrant_client.http import models as qm
from schemas import ChunkRecord
from .ocr import mistral_ocr_pdf
from .preprocess import normalize_arabic, drop_common_headers_footers
from .chuncking import chunk_pages
class ArabicBookRAG:
def __init__(
self,
collection_name: str, # 👈 مهم
embedding_model="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
):
self.collection = collection_name
self.embedder = SentenceTransformer(embedding_model)
self.qdrant = QdrantClient(
url=os.environ["QDRANT_URL"],
api_key=os.environ["QDRANT_API_KEY"],
)
self._ensure_collection()
def _ensure_collection(self):
dim = self.embedder.get_sentence_embedding_dimension()
collections = [c.name for c in self.qdrant.get_collections().collections]
if self.collection not in collections:
self.qdrant.create_collection(
collection_name=self.collection,
vectors_config=qm.VectorParams(size=dim, distance=qm.Distance.COSINE),
)
# =========================
# Ingest PDF (In-Memory)
# =========================
def ingest_pdf(self, pdf_bytes: bytes, meta):
pages = mistral_ocr_pdf(pdf_bytes)
pages = [normalize_arabic(p) for p in pages]
pages = drop_common_headers_footers(pages)
chunks = chunk_pages(pages)
records = [
ChunkRecord(
chunk_id=str(uuid.uuid4()),
doc_id=meta.doc_id,
page_start=ps,
page_end=pe,
text=txt,
title=meta.title,
author=meta.author,
year=meta.year,
)
for txt, ps, pe in chunks
]
vectors = self.embedder.encode(
[r.text for r in records], normalize_embeddings=True
)
self.qdrant.upsert(
collection_name=self.collection,
points=[
qm.PointStruct(
id=r.chunk_id,
vector=v.tolist(),
payload=asdict(r),
)
for r, v in zip(records, vectors)
],
)
return {"pages": len(pages), "chunks": len(records)}
# =========================
# Retrieve
# =========================
def retrieve(
self, queries: List[str], doc_id: Optional[str] = None, top_k: int = 8
):
must = []
if doc_id:
must.append(
qm.FieldCondition(key="doc_id", match=qm.MatchValue(value=doc_id))
)
flt = qm.Filter(must=must) if must else None
hits = []
for q in queries:
qn = normalize_arabic(q)
vec = self.embedder.encode([qn], normalize_embeddings=True)[0]
res = self.qdrant.query_points(
collection_name=self.collection,
query=vec.tolist(),
limit=top_k,
with_payload=True,
query_filter=flt,
).points
hits.extend(res)
return hits
# =========================
# Delete whole book
# =========================
def delete_collection(self):
self.qdrant.delete_collection(self.collection)
|