OntoLearner 1.4.10__py3-none-any.whl → 1.5.0__py3-none-any.whl

ontolearner/learner/rag/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2025 SciKnowOrg
+#
+# Licensed under the MIT License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://opensource.org/licenses/MIT
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .rag import AutoRAGLearner, LLMAugmentedRAGLearner

ontolearner/learner/{rag.py → rag/rag.py}

@@ -14,8 +14,7 @@
 
 import warnings
 from typing import Any
-from ..base import AutoLearner
-
+from ...base import AutoLearner
 
 class AutoRAGLearner(AutoLearner):
     def __init__(self,
@@ -87,3 +86,9 @@ class AutoRAGLearner(AutoLearner):
             return self.llm._non_taxonomic_re_predict(dataset=dataset)
         else:
             warnings.warn("No requirement for fiting the non-taxonomic-re model, the predict module will use the input data to do the fit as well.")
+
+
+class LLMAugmentedRAGLearner(AutoRAGLearner):
+
+    def set_augmenter(self, augmenter):
+        self.retriever.set_augmenter(augmenter=augmenter)

ontolearner/learner/retriever/__init__.py

@@ -16,4 +16,4 @@ from .crossencoder import CrossEncoderRetriever
 from .embedding import GloveRetriever, Word2VecRetriever
 from .ngram import NgramRetriever
 from .learner import AutoRetrieverLearner, LLMAugmentedRetrieverLearner
-from .llm_retriever import LLMAugmenterGenerator, LLMAugmenter, LLMAugmentedRetriever
+from .augmented_retriever import LLMAugmenterGenerator, LLMAugmenter, LLMAugmentedRetriever

ontolearner/learner/retriever/{llm_retriever.py → augmented_retriever.py}

@@ -17,6 +17,8 @@ from typing import Any, List, Dict
 from openai import OpenAI
 import time
 from tqdm import tqdm
+import torch
+import torch.nn.functional as F
 
 from ...base import AutoRetriever
 from ...utils import load_json
@@ -125,7 +127,6 @@ class LLMAugmenterGenerator(ABC):
             except Exception:
                 print("sleep for 5 seconds")
                 time.sleep(5)
-
         return inference
 
     def tasks_data_former(self, data: Any, task: str) -> List[str] | Dict[str, List[str]]:
@@ -298,21 +299,12 @@ class LLMAugmentedRetriever(AutoRetriever):
     Attributes:
         augmenter: An augmenter instance that provides transform() and top_n_candidate.
     """
-
-    def __init__(self) -> None:
-        """
-        Initialize the augmented retriever with no augmenter attached.
-        """
+    def __init__(self, threshold: float = 0.0, cutoff_rate: float = 100.0) -> None:
         super().__init__()
-        self.augmenter = None
+        self.threshold = threshold
+        self.cutoff_rate = cutoff_rate
 
     def set_augmenter(self, augmenter):
-        """
-        Attach an augmenter instance.
-
-        Args:
-            augmenter: An object providing `transform(query, task)` and `top_n_candidate`.
-        """
         self.augmenter = augmenter
 
     def retrieve(self, query: List[str], top_k: int = 5, batch_size: int = -1, task: str = None) -> List[List[str]]:
@@ -328,29 +320,46 @@ class LLMAugmentedRetriever(AutoRetriever):
         Returns:
             list[list[str]]: A list of document lists, one per input query.
         """
-        parent_retrieve = super(LLMAugmentedRetriever, self).retrieve
-
-        if task == 'taxonomy-discovery':
-            query_sets = []
-            for idx in range(self.augmenter.top_n_candidate):
-                query_set = []
-                for qu in query:
-                    query_set.append(self.augmenter.transform(qu, task=task)[idx])
-                query_sets.append(query_set)
-
-            retrieves = [
-                parent_retrieve(query=query_set, top_k=top_k, batch_size=batch_size)
-                for query_set in query_sets
-            ]
-
-            results = []
-            for qu_idx, qu in enumerate(query):
-                qu_result = []
-                for top_idx in range(self.augmenter.top_n_candidate):
-                    qu_result += retrieves[top_idx][qu_idx]
-                results.append(list(set(qu_result)))
-
-            return results
-
-        else:
-            return parent_retrieve(query=query, top_k=top_k, batch_size=batch_size)
+        if task != 'taxonomy-discovery':
+            return super().retrieve(query=query, top_k=top_k, batch_size=batch_size)
+        return self.augmented_retrieve(query, top_k=top_k, batch_size=batch_size, task=task)
+
+    def augmented_retrieve(self, query: List[str], top_k: int = 5, batch_size: int = -1, task: str = None):
+        if self.embeddings is None:
+            raise RuntimeError("Retriever model must index documents before prediction.")
+
+        augmented_queries, index_map = [], []
+        for qu_idx, qu in enumerate(query):
+            augmented = self.augmenter.transform(qu, task=task)
+            for aug in augmented:
+                augmented_queries.append(aug)
+                index_map.append(qu_idx)
+
+        doc_norm = F.normalize(self.embeddings, p=2, dim=1)
+        results = [dict() for _ in range(len(query))]
+
+        if batch_size == -1:
+            batch_size = len(augmented_queries)
+
+        for start in range(0, len(augmented_queries), batch_size):
+            batch_aug = augmented_queries[start:start + batch_size]
+            batch_embeddings = self.embedding_model.encode(batch_aug, convert_to_tensor=True)
+            batch_norm = F.normalize(batch_embeddings, p=2, dim=1)
+            similarity_matrix = torch.matmul(batch_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)
+
+            for i, (doc_indices, sim_scores) in enumerate(zip(topk_indices, topk_similarities)):
+                original_query_idx = index_map[start + i]
+
+                for doc_idx, score in zip(doc_indices.tolist(), sim_scores.tolist()):
+                    if score >= self.threshold:
+                        doc = self.documents[doc_idx]
+                        prev = results[original_query_idx].get(doc, 0.0)
+                        results[original_query_idx][doc] = prev + score
+
+        final_results = []
+        for doc_score_map in results:
+            sorted_docs = sorted(doc_score_map.items(), key=lambda x: x[1], reverse=True)
+            final_results.append([doc for doc, _ in sorted_docs])
+        return final_results

ontolearner/learner/retriever/learner.py

@@ -122,7 +122,6 @@ class AutoRetrieverLearner(AutoLearner):
             warnings.warn("No requirement for fiting the non-taxonomic RE model, the predict module will use the input data to do the fit as well..")
 
 
-
 class LLMAugmentedRetrieverLearner(AutoRetrieverLearner):
 
     def set_augmenter(self, augmenter):
@@ -160,9 +159,9 @@ class LLMAugmentedRetrieverLearner(AutoRetrieverLearner):
             taxonomic_pairs = [{"parent": candidate, "child": query}
                                for query, candidates in zip(data, candidates_lst)
                                for candidate in candidates if candidate.lower() != query.lower()]
-            taxonomic_pairs += [{"parent": query, "child": candidate}
-                                for query, candidates in zip(data, candidates_lst)
-                                for candidate in candidates if candidate.lower() != query.lower()]
+            # taxonomic_pairs += [{"parent": query, "child": candidate}
+            #                    for query, candidates in zip(data, candidates_lst)
+            #                    for candidate in candidates if candidate.lower() != query.lower()]
             unique_taxonomic_pairs, seen = [], set()
             for pair in taxonomic_pairs:
                 key = (pair["parent"].lower(), pair["child"].lower())  # Directional key (parent, child)