diff --git a/CHANGELOG.md b/CHANGELOG.md index 53f81a44..dcc138ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,13 @@ ## Changelog -### v1.4.8 (November 3, 2025) +### v1.4.9 (December 8, 2025) +- add retriever collection +- add documentation for retrievers +- minor bug fixings in docs +- add unittest for retrievers +- add new requirements (`gensim`) + +### v1.4.8 (December 3, 2025) - add alexbeck, rwthdbis, sbunlp, and skhnlp learners - add documentation for learners - minor bug fixings diff --git a/CITATION.cff b/CITATION.cff index 1a7bbe0c..da0e6877 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -31,5 +31,5 @@ keywords: - Large Language Models - Text-to-ontology license: MIT -version: 1.4.7 +version: 1.4.9 date-released: '2025' diff --git a/docs/source/learners/images/Bi_vs_Cross-Encoder.jpg b/docs/source/learners/images/Bi_vs_Cross-Encoder.jpg new file mode 100644 index 00000000..83672d5a Binary files /dev/null and b/docs/source/learners/images/Bi_vs_Cross-Encoder.jpg differ diff --git a/docs/source/learners/images/llm-augmenter.jpg b/docs/source/learners/images/llm-augmenter.jpg new file mode 100644 index 00000000..a4c17c96 Binary files /dev/null and b/docs/source/learners/images/llm-augmenter.jpg differ diff --git a/docs/source/learners/llms4ol_challenge/sbunlp_learner.rst b/docs/source/learners/llms4ol_challenge/sbunlp_learner.rst index c9b7cc0a..860c3a45 100644 --- a/docs/source/learners/llms4ol_challenge/sbunlp_learner.rst +++ b/docs/source/learners/llms4ol_challenge/sbunlp_learner.rst @@ -13,11 +13,11 @@ The team participated in the LLMs4OL 2025 Shared Task, which included four subta .. note:: - Read more about the model at `RWTH-DBIS at LLMs4OL 2024 Tasks A and B Knowledge-Enhanced Domain-Specific Continual Learning and Prompt-Tuning of Large Language Models for Ontology Learning `_. + Read more about the model at `SBU-NLP at LLMs4OL 2025 Tasks A, B, and C: Stage-Wise Ontology Construction Through LLMs Without any Training Procedure `_. .. hint:: - The original implementation is available at `https://github.com/MouYongli/LLMs4OL `_ repository. + The original implementation is available at `https://github.com/rarahnamoun/LLMs4OL-Challenge-ISWC-2025 `_ repository. diff --git a/docs/source/learners/retrieval.rst b/docs/source/learners/retrieval.rst index 9cd39599..5c74e93a 100644 --- a/docs/source/learners/retrieval.rst +++ b/docs/source/learners/retrieval.rst @@ -127,3 +127,257 @@ Similar to LLM learner, Retrieval Learner is also callable via streamlined ``Lea .. hint:: See `Learning Tasks `_ for possible tasks within Learners. + +Customization +----------------------- + +You can easily customize ``AutoRetrieverLearner`` by providing your own base retriever. + +**Example:** + +.. code-block:: python + + from ontolearner.learner import AutoRetrieverLearner + from ontolearner.learner.retriever import NgramRetriever + + # Create a custom retriever (default is AutoRetriever) + retriever_model = NgramRetriever() + + # Use it as the base retriever in the learner + learner = AutoRetrieverLearner(base_retriever=retriever_model) + + # Load a model for retrieval or augmentation + learner.load(model_id='...') + + +.. note:: + + - The ``base_retriever`` must implement the ``AutoRetriever`` interface. + - You can use any compatible retriever, e.g., ``NgramRetriever``, ``Word2VecRetriever``, + or your own custom retriever. + - This allows combining semantic, n-gram, or hybrid retrieval pipelines easily. + + +Retriever Collection +-------------------------- + +NgramRetriever +~~~~~~~~~~~~~~~~~~~~~~~ + + +.. sidebar:: **Supported vectorizers** + + - ``count``: Converts a collection of text documents to a matrix of token counts. + - ``tfidf``: Converts a collection of raw documents to a matrix of TF-IDF features. + + +The ``NgramRetriever`` is a simple, interpretable text retriever based on traditional n-gram vectorization methods, such as `CountVectorizer `_ and `TfidfVectorizer `_. It ranks documents using cosine similarity of n-gram vectors. This is useful for baseline retrieval, keyword matching, or small-scale text search tasks. The following code shows how to import ``NgramRetriever`` and load desired model with desired arguments. + +.. code-block:: python + + from ontolearner.learner import AutoRetrieverLearner + from ontolearner.learner.retriever import NgramRetriever + + retriever = NgramRetriever(ngram_range=(1,2), stop_words='english') + + learner = AutoRetrieverLearner(base_retriever=retriever) + + learner.load(model_id="tfidf") # or "count" + +.. note:: + + For desired arguments refer to `scikit-learn > TfidfVectorizer `_ or `scikit-learn > CountVectorizer `_ + +Word2VecRetriever +~~~~~~~~~~~~~~~~~~~~~~~ + +.. sidebar:: How to Download Word2Vec? + + Download the word2vec from `GoogleNews-vectors-negative300.bin.gz `_ and then you can provide the path inside the ``.load(...)``. + +`Word2Vec `_ retriever encode documents and queries using pre-trained word embeddings. Each document is represented by the average of its word vectors, and retrieval is done via cosine similarity between query vectors and document vectors. The following code shows how to use ``Word2VecRetriever`` inside learner model: + +.. code-block:: python + + from ontolearner.learner import AutoRetrieverLearner + from ontolearner.learner.retriever import Word2VecRetriever + + retriever = Word2VecRetriever() + + learner = AutoRetrieverLearner(base_retriever=retriever) + + learner.load(model_id="path/to/word2vec.bin") # Load pre-trained Word2Vec vectors + +.. note:: + + Learn more about Word2Vec at `https://www.tensorflow.org/text/tutorials/word2vec `_ + +GloveRetriever +~~~~~~~~~~~~~~~~~~~~~~~ +.. sidebar:: How to Download GloVe? + + Download the desired GloVe models from `https://nlp.stanford.edu/projects/glove/ `_ and then you can provide the path inside the ``.load(...)``. + +`GloVe `_ is an unsupervised learning algorithm for obtaining vector representations for words. Training is performed on aggregated global word-word co-occurrence statistics from a corpus, and the resulting representations showcase interesting linear substructures of the word vector space. Here, the ``GloveRetriever`` operates based on GloVe model as shown in the following: + + +.. code-block:: python + + from ontolearner.learner import AutoRetrieverLearner + from ontolearner.learner.retriever import GloveRetriever + + retriever = GloveRetriever() + + learner = AutoRetrieverLearner(base_retriever=retriever) + + learner.load(model_id="path/to/glove.txt") # Load pre-trained GloVe vectors + + +.. hint:: + + In both **Word2Vec** and **GloVe** retrievers, If a word in a word is not in the embedding vocabulary, it is ignored. + +.. note:: + + Refer to the GloVe paper at `GloVe: Global Vectors for Word Representation `_ to learn more about this model. + +CrossEncoderRetriever +~~~~~~~~~~~~~~~~~~~~~~~ + + +.. sidebar:: Cross-Encoder Models + + Collections of publicly available cross-encoder models are available at: `🤗 Sentence Transformers - Cross-Encoders `_. + + +Untill now, the OntoLearner ``AutoRetriever`` (base retriever for ``AutoRetrieverLearner``) were using a Bi-Encoder architecture for retrievals. It is important to understand the difference between Bi- and Cross-Encoder. The following diagram shows the differences: + +.. raw:: html + +
+ Bi-Encoder vs Cross-Encoder +
+
+ + +Bi-Encoders produce for a given sentence a sentence embedding. We pass to a BERT independently the sentences A and B, which result in the sentence embeddings u and v. These sentence embedding can then be compared using cosine similarity. In contrast, for a Cross-Encoder, we pass both sentences simultaneously to the Transformer network. It produces then an output value between 0 and 1 indicating the similarity of the input sentence pair. A Cross-Encoder does not produce a sentence embedding. Also, we are not able to pass individual sentences to a Cross-Encoder (Reference: `Sentence-BERT > Cross-Encoder `_). + + +Here, in the OntoLearner, we implemented a ``CrossEncoderRetriever``, a hybrid dense retriever that combines a BiEncoder for fast candidate retrieval and a CrossEncoder for accurate reranking. Overall ``CrossEncoderRetriever`` uses Bi-Encoder based model for retrieval and Cross-Encoder model for reranking. This provides an efficient and accurate alternative to pure Cross-Encoder or pure Bi-Encoder approaches. To use ``CrossEncoderRetriever`` simply follow the following steps: + + +.. code-block:: python + + from ontolearner.learner import AutoRetrieverLearner + from ontolearner.learner.retriever import CrossEncoderRetriever + + retriever = CrossEncoderRetriever(bi_encoder_model_id='Qwen/Qwen3-Embedding-8B') # pass the bi-encoder model ID used in the first-stage + + learner = AutoRetrieverLearner(base_retriever=retriever) + + learner.load(model_id="cross-encoder/ms-marco-MiniLM-L12-v2") # Model ID for the CrossEncoder (reranking model) here! + # When .load(...) is instantiated, both the bi-encoder and cross-encoder models will be loaded. + + +.. note:: + + Learn more about Retrieve and Rerank approach at `Sentence Transformers > Usage > Retrieve & Re-Rank `_. + +LLMAugmentedRetriever +~~~~~~~~~~~~~~~~~~~~~~~~ +The LLM-Augmented retriever improves retrieval quality by expanding each query into multiple augmented variants using an LLM (e.g., GPT-4). The following diagram shows how LLM-Augmented retriever operates in comparison to usual retriever approach. + + +.. raw:: html + +
+ LLM Augmented Retriever +
+
+ +There are two usage modes: + +**1. Online augmentation (using LLMAugmenterGenerator):** This mode calls the LLM directly to generate augmentation candidates. + +.. code-block:: python + + # Step 1 — Create the generator + from ontolearner.learner.retriever import LLMAugmenterGenerator + llm_augmenter_generator = LLMAugmenterGenerator(model_id='gpt-4.1-mini', token = '...', top_n_candidate=10) + + # Step 2 — Generate augmentations for a dataset + tasks = ['term-typing', 'taxonomy-discovery', 'non-taxonomic-re'] + augments = {"config": llm_augmenter_generator.get_config()} + for task in tasks: + augments[task] = llm_augmenter_generator.augment(data, task=task) + + # Step 3 — Save augmentations + from ontolearner.utils import save_json + save_json("augment.json", augments) + +The online augmentation is designed to avoid multiple calls to the models that may lead into expensive API usage and waiting time. Once the augmenter generator output is stored, it can be used for next stage. + +**2. Offline augmentation (recommended for large experiments):** Instead of calling the LLM repeatedly, you load the previously saved augmentations. + + +.. code-block:: python + + # Step 1 — Load augmenter + from ontolearner.learner.retriever import LLMAugmenter + augmenter = LLMAugmenter("augment.json") + + # Step 2 — Attach it to the retriever + from ontolearner.learner.retriever import LLMAugmentedRetriever + from ontolearner.learner import LLMAugmentedRetrieverLearner + + base_retriever = LLMAugmentedRetriever() + learner = LLMAugmentedRetrieverLearner(base_retriever=base_retriever) + learner.set_augmenter(augmenter) + learner.load(model_id="Qwen/Qwen3-Embedding-8B") # path to desired retriever model. + +Here the ``LLMAugmentedRetrieverLearner`` is the high-level wrapper that orchestrates the loading a retriever model, attaching the ``LLMAugmentedRetriever``, automatically applying LLM-based query expansion during training and prediction, and computing ground truth and returning predictions. + + + +.. list-table:: Summary of Components: + :header-rows: 1 + :widths: 25 75 + + * - Component + - Purpose + * - ``LLMAugmenterGenerator`` + - Calls an LLM (GPT-4, GPT-3.5, etc.) to generate augmentation data. + * - ``LLMAugmenter`` + - Loads offline augmentations (``augment.json``). + * - ``LLMAugmentedRetriever`` + - Expands each query using augmentations before retrieval. + * - ``LLMAugmentedRetrieverLearner`` + - Applies the learner pipeline using the augmented retriever. + +.. rubric:: Example: Using LLMAugmentedRetrieverLearner for Taxonomy Discovery + +.. code-block:: python + + from ontolearner.learner.retriever import LLMAugmenterGenerator, LLMAugmentedRetriever, LLMAugmenter + from ontolearner import LLMAugmentedRetrieverLearner, Wine, train_test_split, evaluation_report + + ontology = Wine() + ontology.load() + ontological_data = ontology.extract() + train_data, test_data = train_test_split(ontological_data, test_size=0.2, random_state=42) + + task="taxonomy-discovery" + + llm_augmenter_generator = LLMAugmenterGenerator(model_id='gpt-4.1-mini', token = 'your_openai_token', top_n_candidate=10) + augments = {"config": llm_augmenter_generator.get_config()} + augments[task] = llm_augmenter_generator.augment(ontological_data, task=task) + + learner.set_augmenter(augments) + learner.load(model_id="Qwen/Qwen3-Embedding-8B") + + # Train, Predict, and Evaluate + learner.fit(train_data, task=task) + predictions = learner.predict(test_data, task=task) + truth = learner.tasks_ground_truth_former(test_data, task=task) + metrics = evaluation_report(truth, predictions, task=task) + print(metrics) diff --git a/ontolearner/VERSION b/ontolearner/VERSION index b2e46d18..4ea2b1f4 100644 --- a/ontolearner/VERSION +++ b/ontolearner/VERSION @@ -1 +1 @@ -1.4.8 +1.4.9 diff --git a/ontolearner/learner/__init__.py b/ontolearner/learner/__init__.py index 0baf580a..f44daab3 100644 --- a/ontolearner/learner/__init__.py +++ b/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 diff --git a/ontolearner/learner/retriever/__init__.py b/ontolearner/learner/retriever/__init__.py new file mode 100644 index 00000000..65c0e037 --- /dev/null +++ b/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 diff --git a/ontolearner/learner/retriever/crossencoder.py b/ontolearner/learner/retriever/crossencoder.py new file mode 100644 index 00000000..7b36766b --- /dev/null +++ b/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 diff --git a/ontolearner/learner/retriever/embedding.py b/ontolearner/learner/retriever/embedding.py new file mode 100644 index 00000000..400b4427 --- /dev/null +++ b/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 diff --git a/ontolearner/learner/retriever.py b/ontolearner/learner/retriever/learner.py similarity index 55% rename from ontolearner/learner/retriever.py rename to ontolearner/learner/retriever/learner.py index 3b5aede5..389e542e 100644 --- a/ontolearner/learner/retriever.py +++ b/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 @@ def _non_taxonomic_re(self, data: Any, test: bool = False) -> Optional[Any]: 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) diff --git a/ontolearner/learner/retriever/llm_retriever.py b/ontolearner/learner/retriever/llm_retriever.py new file mode 100644 index 00000000..0671cc26 --- /dev/null +++ b/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) diff --git a/ontolearner/learner/retriever/ngram.py b/ontolearner/learner/retriever/ngram.py new file mode 100644 index 00000000..bf64ccb1 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 72d4ac1f..76bb5d74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ bitsandbytes="^0.45.1" mistral-common = { version = "^1.8.5", extras = ["sentencepiece"] } protobuf = "<5" Levenshtein = "*" +gensim="*" [tool.poetry.dev-dependencies] ruff = "*" diff --git a/requirements.txt b/requirements.txt index 494f7d29..773523a0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,3 +23,4 @@ bitsandbytes~=0.45.1 mistral-common[sentencepiece]~=1.8.5 protobuf<5 Levenshtein +gensim diff --git a/setup.py b/setup.py index 1dd046bc..ef57aa7c 100644 --- a/setup.py +++ b/setup.py @@ -35,7 +35,8 @@ "scikit-learn>=1.6.1,<2.0.0", "bitsandbytes>=0.45.1,<1.0.0", "protobuf<5", - "Levenshtein" + "Levenshtein", + "gensim" ], classifiers=[ "Development Status :: 5 - Production/Stable", diff --git a/tests/test_cross_encoder_retriever.py b/tests/test_cross_encoder_retriever.py new file mode 100644 index 00000000..850c2b0f --- /dev/null +++ b/tests/test_cross_encoder_retriever.py @@ -0,0 +1,68 @@ +import pytest +import numpy as np +from unittest.mock import MagicMock, patch + +from ontolearner.learner.retriever.crossencoder import CrossEncoderRetriever + + +@pytest.fixture +def mock_models(): + """ + Mock SentenceTransformer, CrossEncoder, and util.semantic_search inside the + ontolearner.learner.retriever.crossencoder module so the retriever can be + tested without loading real models. + """ + module_path = "ontolearner.learner.retriever.crossencoder" + + with patch(f"{module_path}.SentenceTransformer") as MockST, \ + patch(f"{module_path}.CrossEncoder") as MockCE, \ + patch(f"{module_path}.util.semantic_search") as mock_search: + + # Mocked BiEncoder — returns fixed embeddings + mock_bi = MagicMock() + mock_bi.encode.return_value = np.array([[1.0, 0.0], [0.0, 1.0]]) + MockST.return_value = mock_bi + + # Mocked CrossEncoder — deterministic scoring + mock_cross = MagicMock() + mock_cross.predict.return_value = np.array([0.9, 0.1]) # doc1 > doc2 + MockCE.return_value = mock_cross + + # Mocked semantic search — two candidate docs + mock_search.return_value = [ + [{"corpus_id": 0}, {"corpus_id": 1}] + ] + + yield MockST, MockCE, mock_search + + +def test_load(mock_models): + retriever = CrossEncoderRetriever() + retriever.load("test-model") + + assert retriever.bi_encoder is not None + assert retriever.cross_encoder is not None + + +def test_index(mock_models): + retriever = CrossEncoderRetriever() + retriever.load("test-model") + + docs = ["doc1", "doc2"] + retriever.index(docs) + + assert retriever.documents == docs + assert isinstance(retriever.document_embeddings, np.ndarray) + assert retriever.document_embeddings.shape == (2, 2) + + +def test_retrieve(mock_models): + retriever = CrossEncoderRetriever() + retriever.load("test-model") + retriever.index(["doc1", "doc2"]) + + result = retriever.retrieve(["query"], top_k=1) + + # Should return top 1 reranked doc → "doc1" + assert len(result) == 1 + assert result[0] == ["doc1"] diff --git a/tests/test_llm_retriever.py b/tests/test_llm_retriever.py new file mode 100644 index 00000000..3e392549 --- /dev/null +++ b/tests/test_llm_retriever.py @@ -0,0 +1,137 @@ +import pytest +from unittest.mock import MagicMock, patch +from ontolearner.learner.retriever.llm_retriever import ( + LLMAugmenterGenerator, + LLMAugmenter, + LLMAugmentedRetriever, +) +from ontolearner.base import AutoRetriever + + +class DummyData: + """Dummy ontology-like object for testing data formatting.""" + + class TermTyping: + def __init__(self, term): + self.term = term + + class TypingContainer: + def __init__(self, typings): + self.term_typings = typings + + class TaxPair: + def __init__(self, parent, child): + self.parent = parent + self.child = child + + class TaxonomyContainer: + def __init__(self, pairs): + self.taxonomies = pairs + + class NonTax: + def __init__(self, head, tail, relation): + self.head = head + self.tail = tail + self.relation = relation + + class NonTaxContainer: + def __init__(self, triples): + self.non_taxonomies = triples + + +@pytest.fixture +def mock_openai(): + """Patch OpenAI client and return a controlled response for function calling.""" + with patch("ontolearner.learner.retriever.llm_retriever.OpenAI") as mock_client: + instance = mock_client.return_value + + fake_response = MagicMock() + fake_response.choices = [ + MagicMock( + message=MagicMock( + function_call=MagicMock( + arguments="{'candidate_parents': ['A', 'B', 'C', 'D', 'E']}" + ) + ) + ) + ] + + instance.chat.completions.create.return_value = fake_response + yield mock_client + + +def test_tasks_data_former_taxonomy(): + """Test taxonomy data formatting logic.""" + generator = LLMAugmenterGenerator(token="fake") + + data = DummyData() + data.type_taxonomies = DummyData.TaxonomyContainer( + [ + DummyData.TaxPair("Animal", "Dog"), + DummyData.TaxPair("Vehicle", "Car"), + ] + ) + + result = generator.tasks_data_former(data, "taxonomy-discovery") + assert set(result) == {"Animal", "Dog", "Vehicle", "Car"} + + +def test_augment_taxonomy_discovery(mock_openai): + """Ensure LLM function-calling generation works and returns fixed parents.""" + generator = LLMAugmenterGenerator(token="fake", top_n_candidate=5) + + output = generator.augment_taxonomy_discovery(["Dog"]) + assert "Dog" in output + assert output["Dog"] == ["A", "B", "C", "D", "E"] + + +def test_llm_augmenter_transform(): + """Test augmentation lookup behavior.""" + fake_json = { + "config": {"top_n_candidate": 3}, + "taxonomy-discovery": {"Dog": ["Animal", "Mammal", "Pet"]}, + } + + with patch("ontolearner.learner.retriever.llm_retriever.load_json", return_value=fake_json): + augmenter = LLMAugmenter("dummy/path.json") + + assert augmenter.transform("Dog", "taxonomy-discovery") == ["Animal", "Mammal", "Pet"] + assert augmenter.transform("X", "taxonomy-discovery") == ["X"] + + +def test_llm_augmented_retriever_taxonomy(monkeypatch): + retriever = LLMAugmentedRetriever() + + def fake_retrieve(self, query, top_k=5, batch_size=32): + return [[f"doc_{q}_{i}" for i in range(top_k)] for q in query] + + monkeypatch.setattr( + AutoRetriever, + "retrieve", + fake_retrieve + ) + + class FakeAug: + top_n_candidate = 2 + def transform(self, q, task): + return [f"{q}_A", f"{q}_B"] + + retriever.set_augmenter(FakeAug()) + + results = retriever.retrieve(["Dog"], top_k=2, task="taxonomy-discovery") + assert len(results) == 1 + assert len(results[0]) == 4 + + + +def test_llm_augmented_retriever_normal(monkeypatch): + """Test normal retrieval path (no taxonomy discovery).""" + retriever = LLMAugmentedRetriever() + + def fake_retrieve(self, query, top_k=5, batch_size=32): + return [[f"doc_{query[0]}_{i}" for i in range(top_k)]] + + monkeypatch.setattr(AutoRetriever, "retrieve", fake_retrieve) + + results = retriever.retrieve(["Cat"], top_k=3) + assert results == [["doc_Cat_0", "doc_Cat_1", "doc_Cat_2"]] diff --git a/tests/test_ngram_retriever.py b/tests/test_ngram_retriever.py new file mode 100644 index 00000000..1c27c8e3 --- /dev/null +++ b/tests/test_ngram_retriever.py @@ -0,0 +1,40 @@ +import pytest +from ontolearner.learner.retriever.ngram import NgramRetriever + +def test_ngram_retriever_tfidf(): + # Initialize retriever + retriever = NgramRetriever(ngram_range=(1, 2)) + retriever.load("tfidf") + + # Index some documents + docs = ["The quick brown fox", "jumps over the lazy dog", "hello world"] + retriever.index(docs) + + # Check embeddings shape + assert retriever.embeddings.shape[0] == len(docs) + + # Retrieve top 2 documents + results = retriever.retrieve(["quick fox", "hello"], top_k=2) + assert len(results) == 2 + assert all(len(r) <= 2 for r in results) + # Ensure retrieved documents are from original set + for r in results: + for doc in r: + assert doc in docs + +def test_ngram_retriever_count(): + retriever = NgramRetriever(ngram_range=(1, 1)) + retriever.load("count") + + docs = ["apple orange banana", "banana fruit salad", "fruit apple pie"] + retriever.index(docs) + + results = retriever.retrieve(["apple banana"], top_k=2) + assert len(results) == 1 + assert all(doc in docs for doc in results[0]) + +def test_retrieve_without_indexing_raises(): + retriever = NgramRetriever() + retriever.load("tfidf") + with pytest.raises(RuntimeError): + retriever.retrieve(["test query"]) diff --git a/tests/test_word2vec_glove_retrievers.py b/tests/test_word2vec_glove_retrievers.py new file mode 100644 index 00000000..89bd342f --- /dev/null +++ b/tests/test_word2vec_glove_retrievers.py @@ -0,0 +1,90 @@ +import numpy as np +from unittest.mock import MagicMock +import pytest + +from ontolearner.learner.retriever.embedding import ( + Word2VecRetriever, + GloveRetriever, +) + + +# ------------------------- +# Fixtures for mocks +# ------------------------- + +@pytest.fixture +def mock_w2v_model(): + """Mock gensim KeyedVectors for Word2VecRetriever.""" + mock = MagicMock() + mock.vector_size = 3 + mock.__contains__.side_effect = lambda w: w in {"hello", "world"} + mock.__getitem__.side_effect = lambda w: np.array( + [1, 2, 3] if w == "hello" else [4, 5, 6] + ) + return mock + + +@pytest.fixture +def mock_glove_model(): + """Mock GloVe embedding dict.""" + return { + "hello": [1.0, 2.0, 3.0], + "world": [4.0, 5.0, 6.0], + } + + +# ------------------------- +# Word2VecRetriever tests +# ------------------------- + +def test_w2v_index_and_retrieve(mock_w2v_model): + r = Word2VecRetriever() + r.embedding_model = mock_w2v_model # Skip load() + + docs = ["hello world", "hello", "unknown text"] + r.index(docs) + + assert r.embeddings.shape == (3, 3) + + results = r.retrieve(["hello"], top_k=2) + assert len(results) == 1 + assert len(results[0]) == 2 + assert results[0][0] in docs + + +# ------------------------- +# GloveRetriever tests +# ------------------------- + +def test_glove_index_and_retrieve(mock_glove_model): + r = GloveRetriever() + r.embedding_model = mock_glove_model # Skip load() + + docs = ["hello world", "hello"] + r.index(docs) + + assert r.embeddings.shape[0] == 2 + + results = r.retrieve(["world"], top_k=1) + assert len(results) == 1 + assert results[0][0] in docs + + +# ------------------------- +# Encoding tests +# ------------------------- + +def test_w2v_encode_single(mock_w2v_model): + r = Word2VecRetriever() + r.embedding_model = mock_w2v_model + + vec = r._encode_text("hello") + assert np.allclose(vec, np.array([1, 2, 3])) + + +def test_glove_encode_single(mock_glove_model): + r = GloveRetriever() + r.embedding_model = mock_glove_model + + vec = r._encode_text("world") + assert np.allclose(vec, np.array([4, 5, 6]))