OntoLearner 1.4.8__tar.gz → 1.4.10__tar.gz

{ontolearner-1.4.8 → ontolearner-1.4.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: OntoLearner
-Version: 1.4.8
+Version: 1.4.10
 Summary: OntoLearner: A Modular Python Library for Ontology Learning with LLMs.
 License: MIT
 License-File: LICENSE
@@ -17,6 +17,7 @@ Requires-Dist: Levenshtein
 Requires-Dist: bitsandbytes (>=0.45.1,<0.46.0)
 Requires-Dist: dspy (>=2.6.14,<3.0.0)
 Requires-Dist: g4f
+Requires-Dist: gensim
 Requires-Dist: huggingface-hub (>=0.34.4,<0.35.0)
 Requires-Dist: matplotlib
 Requires-Dist: mistral-common[sentencepiece] (>=1.8.5,<2.0.0)
@@ -80,16 +81,16 @@ Please refer to [Installation](https://ontolearner.readthedocs.io/installation.h
 
 ## 🔗 Essential Resources
 
-| Resource                                                                                                                                                                                                            | Info                                                                                                                                                                                                                                                                                                                     |
-|:--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| **[📚 OntoLearner Documentation](https://ontolearner.readthedocs.io/)**                                                                                                                                             | OntoLearner's extensive documentation website.                                                                                                                                                                                                                                                                           |
-| **[🤗 Datasets on Hugging Face](https://huggingface.co/collections/SciKnowOrg/ontolearner-benchmarking-6823bcd051300c210b7ef68a)**                                                                                  | Access curated, machine-readable ontologies.                                                                                                                                                                                                                                                                             |
-| **Quick Tour on OntoLearner** [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1DuElAyEFzd1vtqTjDEXWcc0zCbiV2Yee?usp=sharing) ``version=1.2.1`` | OntoLearner hands-on Colab tutorials. |
-| **[🚀 Quickstart](https://ontolearner.readthedocs.io/quickstart.html)**                                                                                                                                             | Get started quickly with OntoLearner’s main features and workflow.                                                                                                                                                                                                                                                       |
-| **[🕸️ Learning Tasks](https://ontolearner.readthedocs.io/learning_tasks/learning_tasks.html)**                                                                                                                     | Explore supported ontology learning tasks like LLMs4OL Paradigm tasks and Text2Onto.                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                    |
-| **[🧠 Learner Models](https://ontolearner.readthedocs.io/learners/llm.html)**                                                                                                                                       | Browse and configure various learner models, including LLMs, Retrieval, or RAG approaches.                                                                                                                                                                                                                               |
-| **[📚 Ontologies Documentations](https://ontolearner.readthedocs.io/benchmarking/benchmark.html)**                                                                                                                  | Review benchmark ontologies and datasets used for evaluation and training.                                                                                                                                                                                                                                               |
-| **[🧩 How to work with Ontologizer?](https://ontolearner.readthedocs.io/ontologizer/ontology_modularization.html)**                                                                                                 | Learn how to modularize and preprocess ontologies using the Ontologizer module.                                                                                                                                                                                                                                          |
+| Resource                                                                                                                           | Info                                                                                       |
+|:-----------------------------------------------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------|
+| **[📚 OntoLearner Documentation](https://ontolearner.readthedocs.io/)**                                                            | OntoLearner's extensive documentation website.                                             |
+| **[🤗 Datasets on Hugging Face](https://huggingface.co/collections/SciKnowOrg/ontolearner-benchmarking-6823bcd051300c210b7ef68a)** | Access curated, machine-readable ontologies.                                               |
+| **[🚀 Quickstart](https://ontolearner.readthedocs.io/quickstart.html)**                                                            | Get started quickly with OntoLearner’s main features and workflow.                         |
+| **[🕸️ Learning Tasks](https://ontolearner.readthedocs.io/learning_tasks/learning_tasks.html)**                                    | Explore supported ontology learning tasks like LLMs4OL Paradigm tasks and Text2Onto.       |                                                                                                                                                                                                                                    |
+| **[🧠 Learner Models](https://ontolearner.readthedocs.io/learners/llm.html)**                                                      | Browse and configure various learner models, including LLMs, Retrieval, or RAG approaches. |
+| **[📚 Ontologies Documentations](https://ontolearner.readthedocs.io/benchmarking/benchmark.html)**                                 | Review benchmark ontologies and datasets used for evaluation and training.                 |
+| **[🧩 How to work with Ontologizer?](https://ontolearner.readthedocs.io/ontologizer/ontology_modularization.html)**                | Learn how to modularize and preprocess ontologies using the Ontologizer module.            |
+| **[🤗 Ontology Metrics Dashboard](https://huggingface.co/spaces/SciKnowOrg/OntoLearner-Benchmark-Metrics)**                        | Benchmark ontologies with their metrics and complexity scores.                             |
 
 ## 🚀 Quick Tour
 Get started with OntoLearner in just a few lines of code. This guide demonstrates how to initialize ontologies, load datasets, and train an LLM-assisted learner for ontology engineering tasks.
@@ -135,7 +136,7 @@ task = 'non-taxonomic-re'
 ret_learner = AutoRetrieverLearner(top_k=5)
 ret_learner.load(model_id='sentence-transformers/all-MiniLM-L6-v2')
 
-# 5. Fit the model to training data and do the predict
+# 5. Fit the model to training data and then predict over the test data
 ret_learner.fit(train_data, task=task)
 predicts = ret_learner.predict(test_data, task=task)
 

{ontolearner-1.4.8 → ontolearner-1.4.10}/README.md

@@ -39,16 +39,16 @@ Please refer to [Installation](https://ontolearner.readthedocs.io/installation.h
 
 ## 🔗 Essential Resources
 
-| Resource                                                                                                                                                                                                            | Info                                                                                                                                                                                                                                                                                                                     |
-|:--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| **[📚 OntoLearner Documentation](https://ontolearner.readthedocs.io/)**                                                                                                                                             | OntoLearner's extensive documentation website.                                                                                                                                                                                                                                                                           |
-| **[🤗 Datasets on Hugging Face](https://huggingface.co/collections/SciKnowOrg/ontolearner-benchmarking-6823bcd051300c210b7ef68a)**                                                                                  | Access curated, machine-readable ontologies.                                                                                                                                                                                                                                                                             |
-| **Quick Tour on OntoLearner** [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1DuElAyEFzd1vtqTjDEXWcc0zCbiV2Yee?usp=sharing) ``version=1.2.1`` | OntoLearner hands-on Colab tutorials. |
-| **[🚀 Quickstart](https://ontolearner.readthedocs.io/quickstart.html)**                                                                                                                                             | Get started quickly with OntoLearner’s main features and workflow.                                                                                                                                                                                                                                                       |
-| **[🕸️ Learning Tasks](https://ontolearner.readthedocs.io/learning_tasks/learning_tasks.html)**                                                                                                                     | Explore supported ontology learning tasks like LLMs4OL Paradigm tasks and Text2Onto.                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                    |
-| **[🧠 Learner Models](https://ontolearner.readthedocs.io/learners/llm.html)**                                                                                                                                       | Browse and configure various learner models, including LLMs, Retrieval, or RAG approaches.                                                                                                                                                                                                                               |
-| **[📚 Ontologies Documentations](https://ontolearner.readthedocs.io/benchmarking/benchmark.html)**                                                                                                                  | Review benchmark ontologies and datasets used for evaluation and training.                                                                                                                                                                                                                                               |
-| **[🧩 How to work with Ontologizer?](https://ontolearner.readthedocs.io/ontologizer/ontology_modularization.html)**                                                                                                 | Learn how to modularize and preprocess ontologies using the Ontologizer module.                                                                                                                                                                                                                                          |
+| Resource                                                                                                                           | Info                                                                                       |
+|:-----------------------------------------------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------|
+| **[📚 OntoLearner Documentation](https://ontolearner.readthedocs.io/)**                                                            | OntoLearner's extensive documentation website.                                             |
+| **[🤗 Datasets on Hugging Face](https://huggingface.co/collections/SciKnowOrg/ontolearner-benchmarking-6823bcd051300c210b7ef68a)** | Access curated, machine-readable ontologies.                                               |
+| **[🚀 Quickstart](https://ontolearner.readthedocs.io/quickstart.html)**                                                            | Get started quickly with OntoLearner’s main features and workflow.                         |
+| **[🕸️ Learning Tasks](https://ontolearner.readthedocs.io/learning_tasks/learning_tasks.html)**                                    | Explore supported ontology learning tasks like LLMs4OL Paradigm tasks and Text2Onto.       |                                                                                                                                                                                                                                    |
+| **[🧠 Learner Models](https://ontolearner.readthedocs.io/learners/llm.html)**                                                      | Browse and configure various learner models, including LLMs, Retrieval, or RAG approaches. |
+| **[📚 Ontologies Documentations](https://ontolearner.readthedocs.io/benchmarking/benchmark.html)**                                 | Review benchmark ontologies and datasets used for evaluation and training.                 |
+| **[🧩 How to work with Ontologizer?](https://ontolearner.readthedocs.io/ontologizer/ontology_modularization.html)**                | Learn how to modularize and preprocess ontologies using the Ontologizer module.            |
+| **[🤗 Ontology Metrics Dashboard](https://huggingface.co/spaces/SciKnowOrg/OntoLearner-Benchmark-Metrics)**                        | Benchmark ontologies with their metrics and complexity scores.                             |
 
 ## 🚀 Quick Tour
 Get started with OntoLearner in just a few lines of code. This guide demonstrates how to initialize ontologies, load datasets, and train an LLM-assisted learner for ontology engineering tasks.
@@ -94,7 +94,7 @@ task = 'non-taxonomic-re'
 ret_learner = AutoRetrieverLearner(top_k=5)
 ret_learner.load(model_id='sentence-transformers/all-MiniLM-L6-v2')
 
-# 5. Fit the model to training data and do the predict
+# 5. Fit the model to training data and then predict over the test data
 ret_learner.fit(train_data, task=task)
 predicts = ret_learner.predict(test_data, task=task)
 

ontolearner-1.4.10/ontolearner/VERSION

@@ -0,0 +1 @@
+1.4.10

{ontolearner-1.4.8 → ontolearner-1.4.10}/ontolearner/base/ontology.py

@@ -372,7 +372,7 @@ class BaseOntology(ABC):
         # Save updated metrics
         df.to_excel(metrics_file_path, index=False)
 
-    def is_valid_label(label: str) -> Any:
+    def is_valid_label(self, label: str) -> Any:
         invalids = ['root', 'thing']
         if label.lower() in invalids:
             return None
@@ -522,7 +522,7 @@ class BaseOntology(ABC):
                 return True
         return False
 
-    def _is_anonymous_id(label: str) -> bool:
+    def _is_anonymous_id(self, label: str) -> bool:
         """Check if a label represents an anonymous class identifier."""
         if not label:
             return True

{ontolearner-1.4.8 → ontolearner-1.4.10}/ontolearner/learner/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from .llm import AutoLLMLearner, FalconLLM, MistralLLM
-from .retriever import AutoRetrieverLearner
+from .retriever import AutoRetrieverLearner, LLMAugmentedRetrieverLearner
 from .rag import AutoRAGLearner
 from .prompt import StandardizedPrompting
 from .label_mapper import LabelMapper

ontolearner-1.4.10/ontolearner/learner/retriever/__init__.py

@@ -0,0 +1,19 @@
+# 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 .crossencoder import CrossEncoderRetriever
+from .embedding import GloveRetriever, Word2VecRetriever
+from .ngram import NgramRetriever
+from .learner import AutoRetrieverLearner, LLMAugmentedRetrieverLearner
+from .llm_retriever import LLMAugmenterGenerator, LLMAugmenter, LLMAugmentedRetriever

ontolearner-1.4.10/ontolearner/learner/retriever/crossencoder.py

@@ -0,0 +1,129 @@
+# 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.
+import logging
+from typing import List
+from sentence_transformers import CrossEncoder, SentenceTransformer, util
+from tqdm import tqdm
+import numpy as np
+
+from ...base import AutoRetriever
+
+logger = logging.getLogger(__name__)
+
+
+class CrossEncoderRetriever(AutoRetriever):
+    """
+    A hybrid dense retriever that combines a BiEncoder for fast candidate
+    retrieval and a CrossEncoder for accurate reranking.
+
+    This retriever follows a two-stage retrieval process:
+
+    1. **BiEncoder retrieval**:
+       Encodes all documents and queries into embeddings.
+       Computes approximate nearest neighbors to obtain a set of top-k candidates.
+
+    2. **CrossEncoder reranking**:
+       Evaluates each (query, document) pair for semantic relevance.
+       Reranks the initial candidates and outputs the final top results.
+
+    This provides an efficient and accurate alternative to pure CrossEncoder
+    or pure BiEncoder approaches.
+    """
+
+    def __init__(self, bi_encoder_model_id: str = None) -> None:
+        """
+        Initialize the retriever.
+
+        Args:
+            bi_encoder_model_id (str, optional):
+                Model ID for the BiEncoder used in the first-stage retrieval.
+                If not provided, the CrossEncoder model_id passed to `load()`
+                will also be used as the BiEncoder.
+        """
+        super().__init__()
+        self.bi_encoder_model_id = bi_encoder_model_id
+
+    def load(self, model_id: str):
+        """
+        Load both the BiEncoder and CrossEncoder models.
+
+        Args:
+            model_id (str):
+                Model ID for the CrossEncoder (reranking model). If no explicit
+                BiEncoder ID was given at initialization, this ID is also used
+                for the BiEncoder.
+
+        Notes:
+            - BiEncoder is used for fast vector similarity search.
+            - CrossEncoder is used for slow but accurate reranking.
+        """
+        if not self.bi_encoder_model_id:
+            self.bi_encoder_model_id = model_id
+        self.bi_encoder = SentenceTransformer(self.bi_encoder_model_id)
+        self.cross_encoder = CrossEncoder(model_id)
+
+    def index(self, inputs: List[str]):
+        """
+        Pre-encode all documents using the BiEncoder to support efficient
+        semantic search.
+
+        Args:
+            inputs (List[str]):
+                List of documents to index.
+
+        Stores:
+            - `self.documents`: Raw input documents.
+            - `self.document_embeddings`: Tensor of BiEncoder embeddings.
+        """
+        self.documents = inputs
+        self.document_embeddings = self.bi_encoder.encode(inputs, convert_to_tensor=True, show_progress_bar=True)
+
+    def retrieve(self, query: List[str], top_k: int = 5, rerank_k: int = 100, batch_size: int = 32) -> List[List[str]]:
+        """
+        Retrieve top-k most relevant documents per query using a two-stage process.
+
+        Stage 1: Retrieve top `rerank_k` documents using BiEncoder embeddings.
+        Stage 2: Rerank those candidates using the CrossEncoder, returning `top_k`.
+
+        Args:
+            query (List[str]):
+                List of user query strings.
+            top_k (int):
+                Number of final documents to return after reranking.
+            rerank_k (int):
+                Number of candidates to retrieve before reranking.
+            batch_size (int):
+                Batch size for CrossEncoder inference.
+
+        Returns:
+            List[List[str]]:
+                For each query, a list of top-k reranked documents.
+        """
+        results = []
+        # Step 1: Encode queries with the BiEncoder
+        query_embeddings = self.bi_encoder.encode(
+            query, convert_to_tensor=True, show_progress_bar=True
+        )
+        # Step 2: Retrieve candidate documents
+        hits_batch = util.semantic_search(query_embeddings, self.document_embeddings, top_k=rerank_k)
+        # Step 3: Rerank using CrossEncoder
+        for i, hits in enumerate(tqdm(hits_batch, desc="Reranking")):
+            candidates = [self.documents[hit["corpus_id"]] for hit in hits]
+            pairs = [(query[i], doc) for doc in candidates]
+            scores = self.cross_encoder.predict(pairs, batch_size=batch_size, show_progress_bar=False)
+            reranked_idx = np.argsort(scores)[::-1][:top_k]
+            top_docs = [candidates[j] for j in reranked_idx]
+            results.append(top_docs)
+
+        return results

ontolearner-1.4.10/ontolearner/learner/retriever/embedding.py

@@ -0,0 +1,229 @@
+# 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.
+import logging
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+from tqdm import tqdm
+from typing import List, Optional
+from sklearn.metrics.pairwise import cosine_similarity
+from gensim.models import KeyedVectors
+from gensim.utils import simple_preprocess
+
+from ...base import AutoRetriever
+
+logger = logging.getLogger(__name__)
+
+
+class Word2VecRetriever(AutoRetriever):
+    """
+    Retriever that encodes each document by averaging its Word2Vec-style
+    word embeddings. Retrieval is performed by cosine similarity between
+    averaged document vectors and averaged query vectors.
+    """
+
+    def __init__(self) -> None:
+        """
+        Initialize an empty Word2VecRetriever. The model must be loaded using
+        :meth:`load` before indexing or retrieval.
+        """
+        super().__init__()
+        self.embedding_model: Optional[KeyedVectors] = None
+        self.documents: List[str] = []
+        self.embeddings: Optional[torch.Tensor] = None
+
+    def load(self, model_id: str) -> None:
+        """
+        Load a pre-trained Word2Vec KeyedVectors model.
+
+        Args:
+            model_id (str):
+                Path to a Word2Vec `.bin` or `.txt` vector file.
+        """
+        self.embedding_model = KeyedVectors.load_word2vec_format(model_id, binary=True)
+
+    def _encode_text(self, text: str) -> np.ndarray:
+        """
+        Encode text by averaging embeddings for all in-vocabulary words.
+
+        Args:
+            text (str): Input text string.
+
+        Returns:
+            np.ndarray: Averaged embedding vector. If no word is in the vocabulary,
+            a zero vector of appropriate dimensionality is returned.
+        """
+        if self.embedding_model is None:
+            raise RuntimeError("Word2Vec model must be loaded before encoding.")
+
+        words = simple_preprocess(text)
+        valid_vectors = [self.embedding_model[word] for word in words if word in self.embedding_model]
+
+        if not valid_vectors:
+            return np.zeros(self.embedding_model.vector_size)
+
+        return np.mean(valid_vectors, axis=0)
+
+    def index(self, inputs: List[str]) -> None:
+        """
+        Encode and index a list of documents.
+
+        Args:
+            inputs (List[str]): Documents to index.
+
+        Stores:
+            - self.documents: The input documents.
+            - self.embeddings: L2-normalized document embeddings.
+        """
+        self.documents = inputs
+        embeddings = [self._encode_text(doc) for doc in tqdm(inputs)]
+        self.embeddings = F.normalize(torch.tensor(np.stack(embeddings)), p=2, dim=1)
+
+    def retrieve(self, query: List[str], top_k: int = 5, batch_size: int = -1) -> List[List[str]]:
+        """
+        Retrieve the top-k most similar documents for each query.
+
+        Args:
+            query (List[str]): Query texts.
+            top_k (int): Number of results to return per query.
+            batch_size (int): Batch size for processing queries. -1 means all at once.
+
+        Returns:
+            List[List[str]]: One list per query containing top-k matching documents.
+        """
+        if self.embeddings is None:
+            raise RuntimeError("Documents must be indexed before retrieval.")
+
+        query_vec = [self._encode_text(q) for q in query]
+        query_vec = F.normalize(torch.tensor(np.stack(query_vec)), p=2, dim=1)
+
+        if batch_size == -1:
+            batch_size = len(query)
+
+        results = []
+        for i in tqdm(range(0, len(query), batch_size)):
+            q_batch = query_vec[i:i + batch_size]
+            sim = cosine_similarity(q_batch, self.embeddings)
+
+            topk_idx = np.argsort(sim, axis=1)[:, ::-1][:, :top_k]
+
+            for row in topk_idx:
+                results.append([self.documents[j] for j in row])
+
+        return results
+
+
+class GloveRetriever(AutoRetriever):
+    """
+    Retriever that uses GloVe embedding vectors. Each document is encoded
+    by averaging the embeddings of all words that exist in the GloVe vocabulary.
+    """
+
+    def __init__(self) -> None:
+        """
+        Initialize an empty GloveRetriever. Model must be loaded before use.
+        """
+        super().__init__()
+        self.embedding_model: Optional[dict] = None
+        self.documents: List[str] = []
+        self.embeddings: Optional[torch.Tensor] = None
+
+    def load(self, model_id: str) -> None:
+        """
+        Load GloVe embeddings from a text file.
+
+        Args:
+            model_id (str):
+                Path to GloVe `.txt` file, e.g. `glove.6B.300d.txt`.
+        """
+        logger.info(f"Loading GloVe embeddings from {model_id} ...")
+        self.embedding_model = {}
+
+        with open(model_id, "r", encoding="utf8") as f:
+            for line in f:
+                values = line.split()
+                word = values[0]
+                vec = [float(v) for v in values[1:]]
+                self.embedding_model[word] = vec
+
+        logger.info(f"Loaded {len(self.embedding_model)} GloVe words.")
+
+    def _encode_text(self, text: str) -> np.ndarray:
+        """
+        Encode text by averaging GloVe embeddings.
+
+        Args:
+            text (str): Input text.
+
+        Returns:
+            np.ndarray: Averaged embedding vector. Returns zero vector if no words match.
+        """
+        if self.embedding_model is None:
+            raise RuntimeError("GloVe model must be loaded before encoding.")
+
+        words = text.lower().split()
+        vecs = [self.embedding_model[w] for w in words if w in self.embedding_model]
+
+        if not vecs:
+            dim = len(next(iter(self.embedding_model.values())))
+            return np.zeros(dim)
+
+        return np.mean(vecs, axis=0)
+
+    def index(self, inputs: List[str]) -> None:
+        """
+        Index a list of documents by encoding and normalizing them.
+
+        Args:
+            inputs (List[str]): Documents to index.
+        """
+        if self.embedding_model is None:
+            raise RuntimeError("You must load a GloVe model before indexing.")
+
+        self.documents = inputs
+        embeddings = [self._encode_text(doc) for doc in tqdm(inputs)]
+        self.embeddings = F.normalize(torch.tensor(np.stack(embeddings)), p=2, dim=1)
+
+    def retrieve(self, query: List[str], top_k: int = 5, batch_size: int = -1) -> List[List[str]]:
+        """
+        Retrieve top-k most similar documents.
+
+        Args:
+            query (List[str]): Query texts.
+            top_k (int): Number of results per query.
+            batch_size (int): Batch size for query computation.
+
+        Returns:
+            List[List[str]]: Each entry is a list of top-k matching documents.
+        """
+        if self.embeddings is None:
+            raise RuntimeError("Documents must be indexed before retrieval.")
+
+        query_vec = [self._encode_text(q) for q in query]
+        query_vec = F.normalize(torch.tensor(np.stack(query_vec)), p=2, dim=1)
+
+        if batch_size == -1:
+            batch_size = len(query)
+
+        results = []
+        for i in tqdm(range(0, len(query), batch_size)):
+            q_batch = query_vec[i:i + batch_size]
+            sim = cosine_similarity(q_batch, self.embeddings)
+            topk_idx = np.argsort(sim, axis=1)[:, ::-1][:, :top_k]
+
+            for row in topk_idx:
+                results.append([self.documents[j] for j in row])
+
+        return results

ontolearner-1.4.8/ontolearner/learner/retriever.py → ontolearner-1.4.10/ontolearner/learner/retriever/learner.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..base import AutoRetriever, AutoLearner
+from ...base import AutoRetriever, AutoLearner
 from typing import Any, Optional
 import warnings
 
@@ -120,3 +120,98 @@ class AutoRetrieverLearner(AutoLearner):
             return non_taxonomic_re
         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 LLMAugmentedRetrieverLearner(AutoRetrieverLearner):
+
+    def set_augmenter(self, augmenter):
+        self.retriever.set_augmenter(augmenter=augmenter)
+
+    def _retriever_predict(self, data: Any, top_k: int, task: str) -> Any:
+        if isinstance(data, list):
+            return self.retriever.retrieve(query=data, top_k=top_k, batch_size=self._batch_size, task=task)
+        if isinstance(data, str):
+            return self.retriever.retrieve(query=[data], top_k=top_k, task=task)
+        raise TypeError(f"Unsupported data type {type(data)}. You should pass a List[str] or a str.")
+
+    def _term_typing(self, data: Any, test: bool = False) -> Optional[Any]:
+        """
+        during training: data = ["type-1", .... ],
+        during testing: data = ['term-1', ...]
+        """
+        if test:
+            if self._is_term_typing_fit:
+                types = self._retriever_predict(data=data, top_k=self.top_k, task='term-typing')
+                return [{"term": term, "types": type} for term, type in zip(data, types)]
+            else:
+                raise RuntimeError("Term typing model must be fit before prediction.")
+        else:
+            super()._term_typing(data=data, test=test)
+
+    def _taxonomy_discovery(self, data: Any, test: bool = False) -> Optional[Any]:
+        """
+        during training: data = ['type-1', ...],
+        during testing (same data): data= ['type-1', ...]
+        """
+        if test:
+            self._retriever_fit(data=data)
+            candidates_lst = self._retriever_predict(data=data, top_k=self.top_k + 1, task='taxonomy-discovery')
+            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()]
+            unique_taxonomic_pairs, seen = [], set()
+            for pair in taxonomic_pairs:
+                key = (pair["parent"].lower(), pair["child"].lower())  # Directional key (parent, child)
+                if key not in seen:
+                    seen.add(key)
+                    unique_taxonomic_pairs.append(pair)
+            return unique_taxonomic_pairs
+        else:
+            super()._taxonomy_discovery(data=data, test=test)
+
+    def _non_taxonomic_re(self, data: Any, test: bool = False) -> Optional[Any]:
+        """
+        during training: data = ['type-1', ...],
+        during testing: {'types': [...], 'relations': [... ]}
+        """
+        if test:
+            # print(data)
+            if 'types' not in data or 'relations' not in data:
+                raise ValueError("The non-taxonomic re predict should take {'types': [...], 'relations': [... ]}")
+            if len(data['types']) == 0:
+                warnings.warn("No `types` avaliable to do the non-taxonomic re-prediction.")
+                return None
+            self._retriever_fit(data=data['types'])
+            candidates_lst = self._retriever_predict(data=data['types'], top_k=self.top_k + 1, task='non-taxonomic-re')
+            taxonomic_pairs = []
+            taxonomic_pairs_query = []
+            seen = set()
+            for query, candidates in zip(data['types'], candidates_lst):
+                for candidate in candidates:
+                    if candidate != query:
+                        # Directional pair 1: query -> candidate
+                        key1 = (query.lower(), candidate.lower())
+                        if key1 not in seen:
+                            seen.add(key1)
+                            taxonomic_pairs.append((query, candidate))
+                            taxonomic_pairs_query.append(f"Head: {query}\nTail: {candidate}")
+                        # Directional pair 2: candidate -> query
+                        key2 = (candidate.lower(), query.lower())
+                        if key2 not in seen:
+                            seen.add(key2)
+                            taxonomic_pairs.append((candidate, query))
+                            taxonomic_pairs_query.append(f"Head: {candidate}\nTail: {query}")
+
+            self._retriever_fit(data=data['relations'])
+            candidate_relations_lst = self._retriever_predict(data=taxonomic_pairs_query, top_k=self.top_k,
+                                                              task='non-taxonomic-re')
+            non_taxonomic_re = [{"head": head, "tail": tail, "relation": relation}
+                                for (head, tail), candidate_relations in zip(taxonomic_pairs, candidate_relations_lst)
+                                for relation in candidate_relations]
+            return non_taxonomic_re
+        else:
+            super()._non_taxonomic_re(data=data, test=test)