Merge branch 'main' of https://huggingface.co/spaces/yeastcoast/polars-documentation-rag

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,12 @@
+---
+title: Polars Documentation Rag
+emoji: 🐨
+colorFrom: green
+colorTo: blue
+sdk: gradio
+sdk_version: 5.25.2
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

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/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)