OntoLearner 1.4.4__tar.gz → 1.4.5__tar.gz

{ontolearner-1.4.4 → ontolearner-1.4.5}/PKG-INFO

@@ -1,8 +1,9 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: OntoLearner
-Version: 1.4.4
+Version: 1.4.5
 Summary: OntoLearner: A Modular Python Library for Ontology Learning with LLMs.
 License: MIT
+License-File: LICENSE
 Author: Hamed Babaei Giglou
 Author-email: hamedbabaeigiglou@gmail.com
 Requires-Python: >=3.10,<3.14.0

ontolearner-1.4.5/ontolearner/VERSION

@@ -0,0 +1 @@
+1.4.5

{ontolearner-1.4.4 → ontolearner-1.4.5}/ontolearner/base/learner.py

@@ -350,7 +350,7 @@ class AutoRetriever(ABC):
         self.documents = inputs
         self.embeddings = self.embedding_model.encode(inputs, convert_to_tensor=True)
 
-    def retrieve(self, query: List[str], top_k: int = 5) -> List[List[str]]:
+    def retrieve(self, query: List[str], top_k: int = 5, batch_size: int = -1) -> List[List[str]]:
         """
         Retrieve the top-k most similar examples for each query in a list of queries.
 
@@ -363,33 +363,37 @@ class AutoRetriever(ABC):
         """
         if self.embeddings is None:
             raise RuntimeError("Retriever model must index documents before prediction.")
-
-        # Encode all queries at once
         query_embeddings = self.embedding_model.encode(query, convert_to_tensor=True)  # shape: [num_queries, dim]
-
         if query_embeddings.shape[-1] != self.embeddings.shape[-1]:
             raise ValueError(
                 f"Embedding dimension mismatch: query embedding dim={query_embeddings.shape[-1]}, "
                 f"document embedding dim={self.embeddings.shape[-1]}"
             )
-
-        # Normalize embeddings for cosine similarity
-        query_norm = F.normalize(query_embeddings, p=2, dim=1)
         doc_norm = F.normalize(self.embeddings, p=2, dim=1)
+        if batch_size == -1:
+            results = self._retrieve(query_embeddings=query_embeddings, doc_norm=doc_norm, top_k=top_k)
+        else:
+            results = self._batch_retrieve(query_embeddings=query_embeddings, doc_norm=doc_norm, top_k=top_k, batch_size=batch_size)
+        return results
 
-        # Compute cosine similarity: [num_queries, num_docs]
-        similarity_matrix = torch.matmul(query_norm, doc_norm.T)
-
-        # Get top-k indices for each query
-        top_k = min(top_k, len(self.documents))
-        topk_similarities, topk_indices = torch.topk(similarity_matrix, k=top_k, dim=1)
 
-        # Retrieve documents for each query
+    def _retrieve(self, query_embeddings, doc_norm, top_k: int = 5) -> List[List[str]]:
+        query_norm = F.normalize(query_embeddings, p=2, dim=1)
+        similarity_matrix = torch.matmul(query_norm, doc_norm.T)
+        current_top_k = min(top_k, len(self.documents))
+        topk_similarities, topk_indices = torch.topk(similarity_matrix, k=current_top_k, dim=1)
         results = [[self.documents[i] for i in indices] for indices in topk_indices]
-
         return results
 
 
+    def _batch_retrieve(self, query_embeddings, doc_norm, top_k: int = 5, batch_size: int = 1024) -> List[List[str]]:
+        results = []
+        for i in range(0, query_embeddings.size(0), batch_size):
+            batch_queries = query_embeddings[i:i + batch_size]
+            batch_results = self._retrieve(batch_queries, doc_norm, top_k=top_k)
+            results.extend(batch_results)
+        return results
+
 class AutoPrompt(ABC):
     """
     Abstract base class for prompt formatting components.

{ontolearner-1.4.4 → ontolearner-1.4.5}/ontolearner/learner/retriever.py

@@ -17,12 +17,13 @@ from typing import Any, Optional
 import warnings
 
 class AutoRetrieverLearner(AutoLearner):
-    def __init__(self, base_retriever: Any = AutoRetriever(), top_k: int = 5):
+    def __init__(self, base_retriever: Any = AutoRetriever(), top_k: int = 5, batch_size: int = -1):
         super().__init__()
         self.retriever = base_retriever
         self.top_k = top_k
         self._is_term_typing_fit = False
         self._is_taxonomy_discovery_fit = False
+        self._batch_size = batch_size
 
     def load(self, model_id: str = "sentence-transformers/all-MiniLM-L6-v2"):
         self.retriever.load(model_id=model_id)
@@ -35,7 +36,7 @@ class AutoRetrieverLearner(AutoLearner):
 
     def _retriever_predict(self, data:Any, top_k: int) -> Any:
         if isinstance(data, list):
-            return self.retriever.retrieve(query=data, top_k=top_k)
+            return self.retriever.retrieve(query=data, top_k=top_k, batch_size=self._batch_size)
         if isinstance(data, str):
             return self.retriever.retrieve(query=[data], top_k=top_k)
         raise TypeError(f"Unsupported data type {type(data)}. You should pass a List[str] or a str.")

{ontolearner-1.4.4 → ontolearner-1.4.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "OntoLearner"
-version = "1.4.4"
+version = "1.4.5"
 description = "OntoLearner: A Modular Python Library for Ontology Learning with LLMs."
 authors = ["Hamed Babaei Giglou <hamedbabaeigiglou@gmail.com>", "Andrei C. Aioanei <andrei.c.aioanei@gmail.com>"]
 license = "MIT License"

ontolearner-1.4.4/ontolearner/VERSION

@@ -1 +0,0 @@
-1.4.4