langroid 0.1.239__tar.gz → 0.1.240__tar.gz

{langroid-0.1.239 → langroid-0.1.240}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langroid
-Version: 0.1.239
+Version: 0.1.240
 Summary: Harness LLMs with Multi-Agent Programming
 License: MIT
 Author: Prasad Chalasani
@@ -99,8 +99,9 @@ Requires-Dist: sqlalchemy (>=2.0.19,<3.0.0)
 Requires-Dist: tantivy (>=0.21.0,<0.22.0)
 Requires-Dist: thefuzz (>=0.20.0,<0.21.0)
 Requires-Dist: tiktoken (>=0.5.1,<0.6.0)
-Requires-Dist: torch (==2.0.0) ; extra == "hf-embeddings"
+Requires-Dist: torch (==2.0.0) ; extra == "hf-embeddings" or extra == "transformers"
 Requires-Dist: trafilatura (>=1.5.0,<2.0.0)
+Requires-Dist: transformers (>=4.40.1,<5.0.0) ; extra == "transformers"
 Requires-Dist: typer (>=0.9.0,<0.10.0)
 Requires-Dist: types-pillow (>=10.2.0.20240406,<11.0.0.0)
 Requires-Dist: types-pyyaml (>=6.0.12.20240311,<7.0.0.0)

{langroid-0.1.239 → langroid-0.1.240}/langroid/vector_store/qdrantdb.py

@@ -3,7 +3,7 @@ import json
 import logging
 import os
 import uuid
-from typing import List, Optional, Sequence, Tuple, TypeVar
+from typing import Dict, List, Optional, Sequence, Tuple, TypeVar
 
 from dotenv import load_dotenv
 from qdrant_client import QdrantClient
@@ -13,7 +13,12 @@ from qdrant_client.http.models import (
     CollectionStatus,
     Distance,
     Filter,
-    SearchParams,
+    NamedSparseVector,
+    NamedVector,
+    SearchRequest,
+    SparseIndexParams,
+    SparseVector,
+    SparseVectorParams,
     VectorParams,
 )
 
@@ -22,7 +27,7 @@ from langroid.embedding_models.base import (
     EmbeddingModelsConfig,
 )
 from langroid.embedding_models.models import OpenAIEmbeddingsConfig
-from langroid.mytypes import Document, EmbeddingFunction
+from langroid.mytypes import Document, EmbeddingFunction, Embeddings
 from langroid.utils.configuration import settings
 from langroid.vector_store.base import VectorStore, VectorStoreConfig
 
@@ -62,15 +67,36 @@ class QdrantDBConfig(VectorStoreConfig):
     storage_path: str = ".qdrant/data"
     embedding: EmbeddingModelsConfig = OpenAIEmbeddingsConfig()
     distance: str = Distance.COSINE
+    use_sparse_embeddings: bool = False
+    sparse_embedding_model: str = ""
+    sparse_limit: int = 3
 
 
 class QdrantDB(VectorStore):
     def __init__(self, config: QdrantDBConfig = QdrantDBConfig()):
         super().__init__(config)
-        self.config = config
+        self.config: QdrantDBConfig = config
         emb_model = EmbeddingModel.create(config.embedding)
         self.embedding_fn: EmbeddingFunction = emb_model.embedding_fn()
         self.embedding_dim = emb_model.embedding_dims
+        if self.config.use_sparse_embeddings:
+            try:
+                from transformers import AutoModelForMaskedLM, AutoTokenizer
+            except ImportError:
+                raise ImportError(
+                    """
+                    To use sparse embeddings, 
+                    you must install langroid with the [transformers] extra, e.g.:
+                    pip install "langroid[transformers]"
+                    """
+                )
+
+            self.sparse_tokenizer = AutoTokenizer.from_pretrained(
+                self.config.sparse_embedding_model
+            )
+            self.sparse_model = AutoModelForMaskedLM.from_pretrained(
+                self.config.sparse_embedding_model
+            )
         self.host = config.host
         self.port = config.port
         load_dotenv()
@@ -205,12 +231,22 @@ class QdrantDB(VectorStore):
                 else:
                     logger.warning("Recreating fresh collection")
             self.client.delete_collection(collection_name=collection_name)
-        self.client.create_collection(
-            collection_name=collection_name,
-            vectors_config=VectorParams(
+
+        vectors_config = {
+            "": VectorParams(
                 size=self.embedding_dim,
                 distance=Distance.COSINE,
-            ),
+            )
+        }
+        sparse_vectors_config = None
+        if self.config.use_sparse_embeddings:
+            sparse_vectors_config = {
+                "text-sparse": SparseVectorParams(index=SparseIndexParams())
+            }
+        self.client.create_collection(
+            collection_name=collection_name,
+            vectors_config=vectors_config,
+            sparse_vectors_config=sparse_vectors_config,
         )
         collection_info = self.client.get_collection(collection_name=collection_name)
         assert collection_info.status == CollectionStatus.GREEN
@@ -221,6 +257,32 @@ class QdrantDB(VectorStore):
             logger.info(collection_info)
             logger.setLevel(level)
 
+    def get_sparse_embeddings(self, inputs: List[str]) -> List[SparseVector]:
+        if not self.config.use_sparse_embeddings:
+            return []
+        import torch
+
+        tokens = self.sparse_tokenizer(
+            inputs, return_tensors="pt", truncation=True, padding=True
+        )
+        output = self.sparse_model(**tokens)
+        vectors = torch.max(
+            torch.log(torch.relu(output.logits) + torch.tensor(1.0))
+            * tokens.attention_mask.unsqueeze(-1),
+            dim=1,
+        )[0].squeeze(dim=1)
+        sparse_embeddings = []
+        for vec in vectors:
+            cols = vec.nonzero().squeeze().cpu().tolist()
+            weights = vec[cols].cpu().tolist()
+            sparse_embeddings.append(
+                SparseVector(
+                    indices=cols,
+                    values=weights,
+                )
+            )
+        return sparse_embeddings
+
     def add_documents(self, documents: Sequence[Document]) -> None:
         # Add id to metadata if not already present
         super().maybe_add_ids(documents)
@@ -232,6 +294,9 @@ class QdrantDB(VectorStore):
             return
         document_dicts = [doc.dict() for doc in documents]
         embedding_vecs = self.embedding_fn([doc.content for doc in documents])
+        sparse_embedding_vecs = self.get_sparse_embeddings(
+            [doc.content for doc in documents]
+        )
         if self.config.collection_name is None:
             raise ValueError("No collection name set, cannot ingest docs")
         if self.config.collection_name not in colls:
@@ -241,11 +306,16 @@ class QdrantDB(VectorStore):
         # else we get an API error
         b = self.config.batch_size
         for i in range(0, len(ids), b):
+            vectors: Dict[str, Embeddings | List[SparseVector]] = {
+                "": embedding_vecs[i : i + b]
+            }
+            if self.config.use_sparse_embeddings:
+                vectors["text-sparse"] = sparse_embedding_vecs[i : i + b]
             self.client.upsert(
                 collection_name=self.config.collection_name,
                 points=Batch(
                     ids=ids[i : i + b],
-                    vectors=embedding_vecs[i : i + b],
+                    vectors=vectors,
                     payloads=document_dicts[i : i + b],
                 ),
             )
@@ -332,18 +402,39 @@ class QdrantDB(VectorStore):
             filter = Filter()
         else:
             filter = Filter.parse_obj(json.loads(where))
+        requests = [
+            SearchRequest(
+                vector=NamedVector(
+                    name="",
+                    vector=embedding,
+                ),
+                limit=k,
+                with_payload=True,
+                filter=filter,
+            )
+        ]
+        if self.config.use_sparse_embeddings:
+            sparse_embedding = self.get_sparse_embeddings([text])[0]
+            requests.append(
+                SearchRequest(
+                    vector=NamedSparseVector(
+                        name="text-sparse",
+                        vector=sparse_embedding,
+                    ),
+                    limit=self.config.sparse_limit,
+                    with_payload=True,
+                    filter=filter,
+                )
+            )
         if self.config.collection_name is None:
             raise ValueError("No collection name set, cannot search")
-        search_result: List[ScoredPoint] = self.client.search(
-            collection_name=self.config.collection_name,
-            query_vector=embedding,
-            query_filter=filter,
-            limit=k,
-            search_params=SearchParams(
-                hnsw_ef=128,
-                exact=False,  # use Apx NN, not exact NN
-            ),
+        search_result_lists: List[List[ScoredPoint]] = self.client.search_batch(
+            collection_name=self.config.collection_name, requests=requests
         )
+
+        search_result = [
+            match for result in search_result_lists for match in result
+        ]  # 2D list -> 1D list
         scores = [match.score for match in search_result if match is not None]
         docs = [
             Document(**(match.payload))  # type: ignore

{langroid-0.1.239 → langroid-0.1.240}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langroid"
-version = "0.1.239"
+version = "0.1.240"
 description = "Harness LLMs with Multi-Agent Programming"
 authors = ["Prasad Chalasani <pchalasani@gmail.com>"]
 readme = "README.md"
@@ -99,12 +99,13 @@ types-pillow = "^10.2.0.20240406"
 groq = "^0.5.0"
 nest-asyncio = "^1.6.0"
 python-magic = "^0.4.27"
+transformers = {version="^4.40.1", optional=true}
 
 [tool.poetry.extras]
 # install these using `poetry install -E [...]` where [...] is one of the extras below
 # or install multiple extras using, e.g.,  `poetry install -E "litellm mysql"
 hf-embeddings = ["sentence-transformers", "torch"]
-transformers = ["transformers", "huggingface-hub"]
+transformers = ["transformers", "huggingface-hub", "torch"]
 unstructured = ["unstructured"]
 postgres = ["psycopg2", "pytest-postgresql"]
 mysql = ["pymysql", "pytest-mysql"]