OntoLearner 1.4.8__py3-none-any.whl → 1.4.9__py3-none-any.whl

ontolearner/VERSION

@@ -1 +1 @@
-1.4.8
+1.4.9

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/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/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/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/learner/{retriever.py → 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)

ontolearner/learner/retriever/llm_retriever.py

@@ -0,0 +1,356 @@
+# 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 abc import ABC
+from typing import Any, List, Dict
+from openai import OpenAI
+import time
+from tqdm import tqdm
+
+from ...base import AutoRetriever
+from ...utils import load_json
+
+
+class LLMAugmenterGenerator(ABC):
+    """
+    A generator class responsible for creating augmented query candidates using LLMs
+    such as GPT-4 and GPT-3.5. This class provides augmentation support for
+    three ontology-learning tasks:
+
+    - term-typing
+    - taxonomy-discovery
+    - non-taxonomic relation extraction
+
+    For taxonomy discovery, it invokes a function-calling LLM that returns
+    candidate parent classes for each query term.
+
+    Attributes:
+        client (OpenAI): OpenAI API client used for LLM inference.
+        model_id (str): The LLM model identifier.
+        term_typing_function (list): Function call schema for term typing (currently unused).
+        taxonomy_discovery_function (list): Function call schema for taxonomy discovery.
+        non_taxonomic_re_function (list): Function call schema for non-taxonomic relation extraction.
+        top_n_candidate (int): Number of augmented candidates to generate per query.
+        term_typing_prompt (str): Prompt template used for term typing tasks.
+        taxonomy_discovery_prompt (str): Prompt template used for taxonomy discovery.
+        non_taxonomic_re_prompt (str): Prompt template for non-taxonomic RE.
+    """
+
+    def __init__(self, model_id: str = 'gpt-4.1-mini', token: str = '', top_n_candidate: int = 5) -> None:
+        """
+        Initialize the LLM augmenter generator.
+
+        Args:
+            model_id (str): Name of the OpenAI model to use.
+            token (str): API key for authentication.
+            top_n_candidate (int): Number of generated candidate parents per query.
+        """
+        self.client = OpenAI(api_key=token)
+
+        self.model_id = model_id
+
+        self.term_typing_function = []
+        self.taxonomy_discovery_function = [
+            {
+                "name": "discover_taxonomy_parents",
+                "description": "Given a specific type or class (the query), identify potential parent classes that form valid hierarchical (is-a) relationships within a taxonomy.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "candidate_parents": {
+                            "type": "array",
+                            "items": {"type": "string"},
+                            "description": "A ranked list of candidate parent classes representing higher-level categories."
+                        }
+                    },
+                    "required": ["candidate_parents"]
+                }
+            }
+        ]
+
+        self.non_taxonomic_re_function = []
+        self.top_n_candidate = top_n_candidate
+
+        self.term_typing_prompt = ""
+        self.taxonomy_discovery_prompt = (
+            "Given a type (or class) {query}, generate a list of the top {top_n_candidate} candidate classes "
+            "that can form hierarchical (is-a) relationships, where each of these classes is a parent of {query}."
+        )
+        self.non_taxonomic_re_prompt = ""
+
+    def get_config(self) -> Dict[str, Any]:
+        """
+        Get augmenter configuration metadata.
+
+        Returns:
+            dict: Dictionary containing the augmentation configuration.
+        """
+        return {
+            "top_n_candidate": self.top_n_candidate,
+            "augmenter_model": self.model_id
+        }
+
+    def generate(self, conversation, function):
+        """
+        Call an LLM to produce augmented candidates using function-calling.
+
+        Args:
+            conversation (list): Dialogue messages to send to the LLM.
+            function (list): Function schemas supplied to the model.
+
+        Returns:
+            list[str]: A list of top-k generated candidates.
+        """
+        while True:
+            try:
+                completion = self.client.chat.completions.create(
+                    model=self.model_id,
+                    messages=conversation,
+                    functions=function
+                )
+                inference = eval(completion.choices[0].message.function_call.arguments)['candidate_parents'][:self.top_n_candidate]
+                assert len(inference) == self.top_n_candidate
+                break
+            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]]:
+        """
+        Convert raw dataset input into query lists depending on the ontology-learning task.
+
+        Args:
+            data (Any): Input dataset object.
+            task (str): One of {'term-typing', 'taxonomy-discovery', 'non-taxonomic-re'}.
+
+        Returns:
+            List[str] or Dict[str, List[str]]: Formatted query inputs.
+        """
+        formatted_data = []
+        if task == "term-typing":
+            for typing in data.term_typings:
+                formatted_data.append(typing.term)
+            formatted_data = list(set(formatted_data))
+
+        if task == "taxonomy-discovery":
+            for taxonomic_pairs in data.type_taxonomies.taxonomies:
+                formatted_data.append(taxonomic_pairs.parent)
+                formatted_data.append(taxonomic_pairs.child)
+            formatted_data = list(set(formatted_data))
+
+        if task == "non-taxonomic-re":
+            non_taxonomic_types = []
+            non_taxonomic_res = []
+            for triplet in data.type_non_taxonomic_relations.non_taxonomies:
+                non_taxonomic_types.extend([triplet.head, triplet.tail])
+                non_taxonomic_res.append(triplet.relation)
+            formatted_data = {"types": list(set(non_taxonomic_types)), "relations": list(set(non_taxonomic_res))}
+
+        return formatted_data
+
+    def _augment(self, query, conversations, function):
+        """
+        Internal helper to generate augmented candidates for a batch of queries.
+
+        Args:
+            query (list[str]): Input query terms.
+            conversations (list): LLM conversation blocks for each query.
+            function (list): Function-calling schemas.
+
+        Returns:
+            dict[str, list[str]]: Mapping from query → list of augmented candidates.
+        """
+        results = {}
+        for qu, conversation in tqdm(zip(query, conversations)):
+            results[qu] = self.generate(conversation=conversation, function=function)
+        return results
+
+    def augment_term_typing(self, query: List[str]) -> List[str]:
+        """
+        Augment term-typing queries.
+
+        Currently a passthrough: no augmentation is performed.
+
+        Args:
+            query (list[str]): Query terms.
+
+        Returns:
+            list[str]: Unmodified query terms.
+        """
+        return query
+
+    def augment_non_taxonomic_re(self, query: List[str]) -> List[str]:
+        """
+        Augment non-taxonomic relation extraction queries.
+
+        Currently a passthrough.
+
+        Args:
+            query (list[str]): Query terms.
+
+        Returns:
+            list[str]: Unmodified query terms.
+        """
+        return query
+
+    def augment_taxonomy_discovery(self, query: List[str]) -> Dict[str, List[str]]:
+        """
+        Generate augmented candidates for taxonomy discovery.
+
+        Args:
+            query (list[str]): List of type/class names to augment.
+
+        Returns:
+            dict[str, list[str]]: Mapping of original query → list of candidate parents.
+        """
+        conversations = []
+        for qu in query:
+            prompt = self.taxonomy_discovery_prompt.format(query=qu, top_n_candidate=self.top_n_candidate)
+            conversation = [
+                {"role": "system", "content": "Discover possible taxonomy parents."},
+                {"role": "user", "content": prompt}
+            ]
+            conversations.append(conversation)
+
+        return self._augment(query=query, conversations=conversations, function=self.taxonomy_discovery_function)
+
+    def augment(self, data: Any, task: str):
+        """
+        Main entry point for all augmentation modes.
+
+        Args:
+            data (Any): Dataset object to format and augment.
+            task (str): Task type.
+
+        Returns:
+            Any: Augmented output suitable for a retriever.
+
+        Raises:
+            ValueError: If an invalid task type is given.
+        """
+        data = self.tasks_data_former(data=data, task=task)
+        if task == 'term-typing':
+            return self.augment_term_typing(data)
+        elif task == 'taxonomy-discovery':
+            return self.augment_taxonomy_discovery(data)
+        elif task == 'non-taxonomic-re':
+            return self.augment_non_taxonomic_re(data)
+        else:
+            raise ValueError(f"{task} is not a valid task.")
+
+
+class LLMAugmenter:
+    """
+    A lightweight augmenter that loads precomputed augmentation data from disk.
+
+    Attributes:
+        augments (dict): Loaded augmentation data.
+        top_n_candidate (int): Number of augmentation candidates per query.
+    """
+
+    def __init__(self, path: str) -> None:
+        """
+        Initialize an augmenter that uses offline augmentation data.
+
+        Args:
+            path (str): Path to a JSON file containing saved augmentations.
+        """
+        self.augments = load_json(path)
+        self.top_n_candidate = self.augments['config']['top_n_candidate']
+
+    def transform(self, query: str, task: str) -> List[str]:
+        """
+        Retrieve the augmented versions of a query term for a specific task.
+
+        Args:
+            query (str): Input query term.
+            task (str): Task identifier.
+
+        Returns:
+            list[str]: Augmented query candidates.
+        """
+        if task == 'taxonomy-discovery':
+            return self.augments[task].get(query, [query])
+        else:
+            return [query]
+
+
+class LLMAugmentedRetriever(AutoRetriever):
+    """
+    A retriever that enhances queries using LLM-based augmentation before retrieving documents.
+
+    Supports special augmentation logic for taxonomy discovery where each input query
+    is expanded into several augmented variants.
+
+    Attributes:
+        augmenter: An augmenter instance that provides transform() and top_n_candidate.
+    """
+
+    def __init__(self) -> None:
+        """
+        Initialize the augmented retriever with no augmenter attached.
+        """
+        super().__init__()
+        self.augmenter = None
+
+    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]]:
+        """
+        Retrieve documents for a batch of queries, optionally using query augmentation.
+
+        Args:
+            query (list[str]): List of input query terms.
+            top_k (int): Number of documents to retrieve.
+            batch_size (int): Batch size for retrieval.
+            task (str): Optional task identifier that determines augmentation behavior.
+
+        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)

ontolearner/learner/retriever/ngram.py

@@ -0,0 +1,123 @@
+# 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 numpy as np
+from typing import List
+from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+from tqdm import tqdm
+
+from ...base import AutoRetriever
+
+logger = logging.getLogger(__name__)
+
+
+class NgramRetriever(AutoRetriever):
+    """
+    A retriever based on traditional n-gram vectorization methods such as TF-IDF
+    and CountVectorizer.
+
+    This retriever converts documents and queries into sparse bag-of-ngrams
+    vectors and ranks documents using cosine similarity. It is simple,
+    interpretable, and suitable for small-scale baselines or non-semantic
+    text matching.
+    """
+
+    def __init__(self, **vectorizer_kwargs) -> None:
+        """
+        Initialize the n-gram retriever.
+
+        Args:
+            **vectorizer_kwargs: Additional keyword arguments passed directly
+                to the scikit-learn vectorizer (e.g., ngram_range, stop_words).
+        """
+        super().__init__()
+        self.vectorizer_kwargs = vectorizer_kwargs
+        self.vectorizer = None
+        self.embeddings = None
+
+    def load(self, model_id) -> None:
+        """
+        Load and initialize the vectorizer based on `model_id`.
+
+        Args:
+            model_id (str): Either `"tfidf"` for TF-IDF or `"count"` for
+                CountVectorizer.
+
+        Raises:
+            ValueError: If the model_id is not one of the supported options.
+        """
+        if model_id == "tfidf":
+            self.vectorizer = TfidfVectorizer(**self.vectorizer_kwargs)
+        elif model_id == "count":
+            self.vectorizer = CountVectorizer(**self.vectorizer_kwargs)
+        else:
+            raise ValueError(f"Invalid mode '{model_id}'. Choose from ['tfidf', 'count'].")
+
+    def index(self, inputs: List[str]) -> None:
+        """
+        Fit the vectorizer and index (vectorize) the input documents.
+
+        Args:
+            inputs (List[str]): List of text documents to index.
+
+        Notes:
+            This method must be run before calling `retrieve()`. It creates the
+            document embedding matrix used for similarity search.
+        """
+        if self.vectorizer is None:
+            # Default to TF-IDF if the user never called `load()`
+            self.load(model_id="tfidf")
+
+        self.documents = inputs
+        logger.info("Fitting vectorizer and transforming documents...")
+        self.embeddings = self.vectorizer.fit_transform(inputs)
+        logger.info(f"Document embeddings created with shape: {self.embeddings.shape}")
+
+    def retrieve(self, query: List[str], top_k: int = 5, batch_size: int = -1) -> List[List[str]]:
+        """
+        Retrieve the most similar documents for each query string.
+
+        Args:
+            query (List[str]): A list of query strings.
+            top_k (int): Number of most similar documents to return per query.
+            batch_size (int): Number of queries to process at once.
+                Use `-1` to process all queries in a single batch.
+
+        Returns:
+            List[List[str]]: For each query, a list containing the top-k
+            matching documents.
+
+        Raises:
+            RuntimeError: If retrieval is attempted before indexing.
+        """
+        if self.embeddings is None:
+            raise RuntimeError("Retriever must index documents before calling `retrieve()`.")
+
+        logger.info("Vectorizing query text...")
+        query_vec = self.vectorizer.transform(query)
+        logger.info(f"Query vectors created with shape: {query_vec.shape}")
+
+        results = []
+        if batch_size == -1:
+            batch_size = len(query)
+
+        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_indices in topk_idx:
+                results.append([self.documents[j] for j in row_indices])
+
+        return results

{ontolearner-1.4.8.dist-info → ontolearner-1.4.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: OntoLearner
-Version: 1.4.8
+Version: 1.4.9
 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.dist-info → ontolearner-1.4.9.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ontolearner/VERSION,sha256=ffNZhUJ2DAagaoMtfLwSvSCz_Rvq5cjAOZ6WQvcHNxc,6
+ontolearner/VERSION,sha256=x-xbkXEIv48hifmVFcVtJDdZj6d_bmXwy3Lp4d5pPVY,6
 ontolearner/__init__.py,sha256=E4yukFv2PV4uyztTPDWljCySY9AVDcDDzabuvxfabYE,1889
 ontolearner/_learner.py,sha256=2CRQvpsz8akIOdxTs2-KLJ-MssULrjpK-QDD3QXUJXI,5297
 ontolearner/_ontology.py,sha256=W1mp195SImqLKwaj4ueEaBWuLJg2jUdx1JT20Ds3fmQ,6950
@@ -12,12 +12,17 @@ ontolearner/data_structure/metric.py,sha256=4QKkZ5L1YK6hDTU-N5Z9I9Ha99DVHmGfYxK7
 ontolearner/evaluation/__init__.py,sha256=4BZr3BUXjQDTj4Aqlqy4THa80lZPsMuh1EBTCyi9Wig,842
 ontolearner/evaluation/evaluate.py,sha256=NYCVcmPqpyIxYZrMAim37gL-erdh698RD3t3eNTTgZc,1163
 ontolearner/evaluation/metrics.py,sha256=3Aw6ycJ3_Q6xfj4tMBJP6QcexUei0G16H0ZQWt87aRU,6286
-ontolearner/learner/__init__.py,sha256=L54a3uvOeS6hbID6-BFd6fo9rH4WK2Q1XecpXMkEru0,768
+ontolearner/learner/__init__.py,sha256=RKREPrrjzQ5KYvcOwC_2l7yFKwFBd6HoCwhX2H6Spg8,798
 ontolearner/learner/label_mapper.py,sha256=YMPeFKzJxoCYNU5z7QRYPbB88sWdu1iT6iBDpPsjn-4,3792
 ontolearner/learner/llm.py,sha256=3kq_IrwEPTFgeNVKZH9Er_OydJuDpRBtM3YXNNa8_KA,10343
 ontolearner/learner/prompt.py,sha256=0ckH7xphIDKczPe7G-rwiOxFGZ7RsLnpPlNW92b-31U,1574
 ontolearner/learner/rag.py,sha256=eysB2RvcWkVo53s8-kSbZtJv904YVTmdtxplM4ukUKM,4283
-ontolearner/learner/retriever.py,sha256=PNDAwsLIOBD3aQW2Ez0q6PqE3CB7d_GN-yLKJ9_D04s,6204
+ontolearner/learner/retriever/__init__.py,sha256=G5XuJcTblqXVWboVW9StJ2Vo2xACp_kG5_w2nrueqlc,854
+ontolearner/learner/retriever/crossencoder.py,sha256=yurzGE4zydlBSwUefi1CugsWv34HEZ61qADG_-nILbo,4996
+ontolearner/learner/retriever/embedding.py,sha256=Lp9oA7LiOYaSWDvzG779KMv5keNl6Xv7hw0WpeaepDE,7875
+ontolearner/learner/retriever/learner.py,sha256=VcarTwwR8HNddJCh0loCQejDzZ_GO4NkdQUjEhLVy48,11181
+ontolearner/learner/retriever/llm_retriever.py,sha256=goInWYxrD9PSo_EsSKbNV8wEaSPvWY3LEC8XM7jlH64,12917
+ontolearner/learner/retriever/ngram.py,sha256=XgS1OeheKEIi7wfJHZgS8mWxKv9MQrP0apOJD_XSOnM,4575
 ontolearner/learner/taxonomy_discovery/__init__.py,sha256=-Hb5Dl6_6c4l1uIT2zWtyBWMq5cjVD4PNjxt5qJePl4,747
 ontolearner/learner/taxonomy_discovery/alexbek.py,sha256=kFEDvoKxLf-sB7-d5REkcC0DqXZpcA6ZSJ2QHrNoC5E,19010
 ontolearner/learner/taxonomy_discovery/rwthdbis.py,sha256=698Gze2cR-QIhpTbuaOFm7Q4p0lCbdWz3rO6rewJZ1s,41644
@@ -65,7 +70,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.8.dist-info/METADATA,sha256=B_ULVAw849kBqCpF0-oX8lUuq2d7GFMLPPESYQdWnp8,14158
-ontolearner-1.4.8.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-ontolearner-1.4.8.dist-info/licenses/LICENSE,sha256=krXMLuMKgzX-UgaufgfJdm9ojIloZot7ZdvJUnNxl4I,1067
-ontolearner-1.4.8.dist-info/RECORD,,
+ontolearner-1.4.9.dist-info/METADATA,sha256=c_V_1mUkxAhzJz04u1wRYU7xodpZQdiJXBVFzUCIMK8,11444
+ontolearner-1.4.9.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+ontolearner-1.4.9.dist-info/licenses/LICENSE,sha256=krXMLuMKgzX-UgaufgfJdm9ojIloZot7ZdvJUnNxl4I,1067
+ontolearner-1.4.9.dist-info/RECORD,,