Changes from ggruber193/polars-docu-chat-rag

requirements.txt

@@ -0,0 +1,7 @@
+beautifulsoup4~=4.13.4
+markdown~=3.8
+langchain~=0.3.23
+transformers~=4.51.3
+torch~=2.6.0
+tqdm~=4.67.1
+qdrant_client

src/config.py

@@ -0,0 +1,12 @@
+EMBEDDING_MODEL = "thenlper/gte-small"
+
+QDRANT_COLLECTION_NAME = "polars-documentation"
+
+def get_qdrant_config():
+    from qdrant_client import models
+    QDRANT_COLLECTION_CONFIG = {
+        "collection_name": QDRANT_COLLECTION_NAME,
+        "vectors_config": models.VectorParams(size=384, distance=models.Distance.COSINE),   # on_disk=True),
+        # "hnsw_config": models.HnswConfigDiff(on_disk=True)
+    }
+    return QDRANT_COLLECTION_CONFIG

src/data_processing/embeddings.py

@@ -0,0 +1,37 @@
+from transformers import AutoModel, AutoTokenizer
+from torch import Tensor
+from torch import functional as F
+
+from src.config import EMBEDDING_MODEL
+from src.utils import batched
+
+
+class TextEmbedder:
+    def __init__(self, modelname=EMBEDDING_MODEL, max_length=512):
+        self.tokenizer = AutoTokenizer.from_pretrained(modelname)
+        self.model = AutoModel.from_pretrained(modelname)
+        self.max_length = max_length
+
+    @staticmethod
+    def average_pool(last_hidden_states: Tensor,
+                     attention_mask: Tensor) -> Tensor:
+        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+    def embed_text(self, text: str | list[str], batch_size=128):
+        if isinstance(text, str):
+            text = [text]
+
+        outputs = []
+
+        for batch in batched(text, n=batch_size):
+            batch_dict = self.tokenizer(batch, max_length=self.max_length, padding=True, truncation=True, return_tensors='pt')
+            output = self.model(**batch_dict)
+            embeddings = self.average_pool(output.last_hidden_state, batch_dict['attention_mask'])
+
+            # embeddings = F.norm(embeddings, p=2, dim=1)
+            # scores = (embeddings[:1] @ embeddings[1:].T) * 100
+
+            embeddings = embeddings.tolist()
+            outputs += embeddings
+        return outputs

src/data_processing/process_markdown.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from bs4 import BeautifulSoup
+from langchain_core.documents import Document
+from markdown import markdown
+from pathlib import Path
+from langchain.text_splitter import MarkdownTextSplitter, MarkdownHeaderTextSplitter, TextSplitter
+
+from src.utils import batched
+
+
+def read_markdown_file(path: str | Path) -> [str, str]:
+    path = Path(path)
+    with open(path, 'r', encoding="utf8") as f_r:
+        text = f_r.read()
+
+    # text = markdown(text)
+    # text = ''.join(BeautifulSoup(text).findAll(text=True))
+    return text, str(path)
+
+
+def split_markdown(md: str | list[str],
+                   metadata=dict[str, Any] | list[dict[str, Any]],
+                   chunk_size=512,
+                   overlap=64,
+                   splitter: TextSplitter = None) -> list[Document]:
+    if isinstance(md, str):
+        md = [md]
+        if isinstance(metadata, list):
+            raise ValueError("metadata should be a single dict")
+        metadata = [metadata]
+    if splitter is None:
+        headers_to_split_on = [
+            ("#", "Header 1"),
+            ("##", "Header 2"),
+            ("###", "Header 3"),
+        ]
+        md = [MarkdownHeaderTextSplitter(headers_to_split_on, strip_headers=False).split_text(i) for i in md]
+        metadata = [{**metadata[i], **text.metadata} for i, text_split in enumerate(md) for text in text_split]
+        md = [j.page_content for i in md for j in i]
+        splitter = MarkdownTextSplitter(chunk_size=chunk_size, chunk_overlap=overlap)
+
+    docs = splitter.create_documents(md, metadata)
+    return docs
+
+
+def process_markdown_files(paths: list[str | Path], batch_size=1, chunk_size=512, overlap=64):
+    for files in batched(paths, batch_size):
+        mds_w_paths = [read_markdown_file(i) for i in files]
+        metadata = [{"path": md_path} for _, md_path in mds_w_paths]
+        md = [md for md, _ in mds_w_paths]
+        docs = split_markdown(md, metadata, chunk_size=chunk_size, overlap=overlap)
+        yield [i.page_content for i in docs], [i.metadata for i in docs]

src/data_processing/upload_to_qdrant.py

@@ -0,0 +1,61 @@
+from typing import Any
+
+from qdrant_client import QdrantClient, models
+from uuid import uuid4
+
+from src.config import QDRANT_COLLECTION_NAME
+
+
+class QdrantStore:
+    def __init__(self, client: QdrantClient, collection_config=None):
+        self.client = client
+        self.collection_names = set([i.name for i in client.get_collections().collections])
+
+        if collection_config is not None:
+            self.create_collection(collection_config)
+
+    def create_collection(self, collection_config: dict):
+        collection_name = collection_config["collection_name"]
+        if not self.client.collection_exists(collection_name):
+            self.client.create_collection(**collection_config)
+            self.collection_names.add(collection_name)
+
+    def _check_collection_name(self, collection_name):
+        if collection_name not in self.collection_names:
+            raise ValueError(f"Collection: {collection_name} does not exist.")
+
+    def upsert_points(self,
+                      vectors: Any | list[Any],
+                      payloads: dict | list[dict],
+                      collection_name: str):
+        self._check_collection_name(collection_name)
+
+        ids = [str(uuid4()) for _ in payloads]
+
+        self.client.upsert(
+            collection_name=collection_name,
+            points=models.Batch(
+                ids=ids,
+                payloads=payloads,
+                vectors=vectors
+            )
+        )
+
+    def delete_points(self,
+                      filters: dict[str, list[models.FieldCondition]],
+                      collection_name: str):
+        self._check_collection_name(collection_name)
+
+        self.client.delete(
+            collection_name=collection_name,
+            points_selector=models.Filter(**filters)
+        )
+
+    def delete_points_by_match(self,
+                               key_value: tuple[str, list[str] | str],
+                               collection_name: str):
+        key, values = key_value
+        if isinstance(values, str):
+            values = [values]
+        filter = {"must": [models.FieldCondition(key=key, match=models.MatchAny(any=values))]}
+        self.delete_points(filter, collection_name)

src/testing.py

@@ -0,0 +1,19 @@
+from bs4 import BeautifulSoup
+from markdown import markdown
+from langchain.text_splitter import MarkdownTextSplitter
+
+
+path = "D:\PycharmProjects\polargs-docu-chat-rag\data\polars-docu\concepts\data-types-and-structures.md"
+
+with open(path, 'r', encoding="utf8") as f_r:
+    test_md = f_r.read()
+
+html = markdown(test_md)
+text = ''.join(BeautifulSoup(html).findAll(text=True))
+
+print(text[:10])
+
+splitter = MarkdownTextSplitter(chunk_size=512, chunk_overlap=64)
+
+docs = splitter.create_documents([text])
+print(docs)

src/utils.py

@@ -0,0 +1,12 @@
+from itertools import islice
+
+
+def batched(iterable, n, *, strict=False):
+    # batched('ABCDEFG', 3) → ABC DEF G
+    if n < 1:
+        raise ValueError('n must be at least one')
+    iterator = iter(iterable)
+    while batch := tuple(islice(iterator, n)):
+        if strict and len(batch) != n:
+            raise ValueError('batched(): incomplete batch')
+        yield batch