OntoLearner 1.4.4__py3-none-any.whl → 1.4.5__py3-none-any.whl

ontolearner/VERSION

@@ -1 +1 @@
-1.4.4
+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/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.dist-info → ontolearner-1.4.5.dist-info}/METADATA

@@ -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.4.dist-info → ontolearner-1.4.5.dist-info}/RECORD

@@ -1,9 +1,9 @@
-ontolearner/VERSION,sha256=0D2LJotBcTKKwkWM6MqC1pof3jhGnU5PquT4-o2KXjU,6
+ontolearner/VERSION,sha256=9DoCksZ0VCotZ2oPT1iKf_vhy6wlCQ3zRLEv8SefeCA,6
 ontolearner/__init__.py,sha256=E4yukFv2PV4uyztTPDWljCySY9AVDcDDzabuvxfabYE,1889
 ontolearner/_learner.py,sha256=2CRQvpsz8akIOdxTs2-KLJ-MssULrjpK-QDD3QXUJXI,5297
 ontolearner/_ontology.py,sha256=W1mp195SImqLKwaj4ueEaBWuLJg2jUdx1JT20Ds3fmQ,6950
 ontolearner/base/__init__.py,sha256=5pf-ltxzGp32xhEcPdbtm11wXJrYJMUeWG-mbcAYD8Q,705
-ontolearner/base/learner.py,sha256=DVWp7OHlhTYU3Es7Q6CWCOeL7Y5LTbjWilTri_DNExs,17897
+ontolearner/base/learner.py,sha256=6iQlrfffK6RtEZKA2NvXExQEA7hVvlyDg4zWM0fkaCQ,18497
 ontolearner/base/ontology.py,sha256=JbMJ1-WUyHWQiNJL-DeaqcriUimLdqN3_ESROgqOPTQ,24772
 ontolearner/base/text2onto.py,sha256=iUXYZoqnwgebQuQzM-XSGTVRfHLlhjUK_z5XUvhRICc,5388
 ontolearner/data_structure/__init__.py,sha256=1HiKvk8FKjhYeI92RHnJXxyQbUJBi3JFytjQjthsY_s,599
@@ -17,7 +17,7 @@ ontolearner/learner/label_mapper.py,sha256=-XW8MHafm4ix3e9u-RRwDePJ71D804DNuKzdf
 ontolearner/learner/llm.py,sha256=bwCoeR7z3YgYrkKyjDM-MRHZAuDzpUt8f-A0bDUbtGM,7151
 ontolearner/learner/prompt.py,sha256=0ckH7xphIDKczPe7G-rwiOxFGZ7RsLnpPlNW92b-31U,1574
 ontolearner/learner/rag.py,sha256=eysB2RvcWkVo53s8-kSbZtJv904YVTmdtxplM4ukUKM,4283
-ontolearner/learner/retriever.py,sha256=FIsvutDXvrr9N6AMu35TNJHdiQGbmRQ4TTGfRRdHdYo,4931
+ontolearner/learner/retriever.py,sha256=8X4fVD-OJCMLPH5Tl3nUOuXVhPLApzLAJcgIY5p9gEI,5020
 ontolearner/ontology/__init__.py,sha256=F9Ta1qCX9mOxIK5CPRypEoglQNkpJ6SJpqziz73xKQE,1328
 ontolearner/ontology/agriculture.py,sha256=ZaXHNEFjbtsMH8M7HQ8ypnfJS4TUQy_as16fwv-kOKA,5903
 ontolearner/ontology/arts_humanities.py,sha256=K4ceDJL6PfIfSJZ86uQUkUXOVoiERG6ItgvVE2lhLKk,3996
@@ -53,7 +53,7 @@ ontolearner/tools/visualizer.py,sha256=cwijl4yYaS1SCLM5wbvRTEcbQj9Bjo4fHzZR6q6o8
 ontolearner/utils/__init__.py,sha256=pSEyU3dlPMADBqygqaaid44RdWf0Lo3Fvz-K_rQ7_Bw,733
 ontolearner/utils/io.py,sha256=3DqGK2p7c0onKi0Xxs16WB08uHfHUId3bW0dDKwyS0g,2110
 ontolearner/utils/train_test_split.py,sha256=Zlm42eT6QGWwlySyomCPIiTGmGqeN_h4z4xBY2EAOR8,11530
-ontolearner-1.4.4.dist-info/LICENSE,sha256=krXMLuMKgzX-UgaufgfJdm9ojIloZot7ZdvJUnNxl4I,1067
-ontolearner-1.4.4.dist-info/METADATA,sha256=Uf9twY6zgfxNZbgmtCqIhYJ1nnzzwe-TzC5_bztEV_U,13999
-ontolearner-1.4.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-ontolearner-1.4.4.dist-info/RECORD,,
+ontolearner-1.4.5.dist-info/METADATA,sha256=89wwW4YjejZ9_TdDjn5JwcCQLK-Khb_DvDrBdRGFeJ8,14021
+ontolearner-1.4.5.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
+ontolearner-1.4.5.dist-info/licenses/LICENSE,sha256=krXMLuMKgzX-UgaufgfJdm9ojIloZot7ZdvJUnNxl4I,1067
+ontolearner-1.4.5.dist-info/RECORD,,

{ontolearner-1.4.4.dist-info → ontolearner-1.4.5.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.1.3
+Generator: poetry-core 2.2.0
 Root-Is-Purelib: true
 Tag: py3-none-any