diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2d4a9abafc..9eb41d1754 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -798,4 +798,4 @@ jobs: with: status: ${{ job.status }} channel: '#haystack' - if: failure() && github.repository_owner == 'deepset-ai' && github.ref == 'refs/heads/main' + if: failure() && github.repository_owner == 'deepset-ai' && github.ref == 'refs/heads/main' \ No newline at end of file diff --git a/docs/_src/api/api/primitives.md b/docs/_src/api/api/primitives.md index 28d60acd96..604a15c76c 100644 --- a/docs/_src/api/api/primitives.md +++ b/docs/_src/api/api/primitives.md @@ -2,6 +2,12 @@ # Module schema + + +#### ContentTypes + +Types of supported content_types. + ## Document @@ -17,7 +23,7 @@ class Document() ```python def __init__(content: Union[str, pd.DataFrame], - content_type: Literal["text", "table", "image", "audio"] = "text", + content_type: ContentTypes = "text", id: Optional[str] = None, score: Optional[float] = None, meta: Optional[Dict[str, Any]] = None, @@ -38,7 +44,7 @@ There's an easy option to convert from/to dicts via `from_dict()` and `to_dict`. **Arguments**: - `content`: Content of the document. For most cases, this will be text, but it can be a table or image. -- `content_type`: One of "text", "table" or "image". Haystack components can use this to adjust their +- `content_type`: One of "text", "table", "image" or "audio". Haystack components can use this to adjust their handling of Documents and check compatibility. - `id`: Unique ID for the document. If not supplied by the user, we'll generate one automatically by creating a hash from the supplied text. This behaviour can be further adjusted by `id_hash_keys`. diff --git a/haystack/document_stores/memory.py b/haystack/document_stores/memory.py index 4840dc8fe4..9e0a0f625f 100644 --- a/haystack/document_stores/memory.py +++ b/haystack/document_stores/memory.py @@ -383,7 +383,9 @@ class InMemoryDocumentStore(BaseDocumentStore): candidate_docs = [] for doc, score in zip(document_to_search, scores): curr_meta = deepcopy(doc.meta) - new_document = Document(id=doc.id, content=doc.content, meta=curr_meta, embedding=doc.embedding) + new_document = Document( + id=doc.id, content=doc.content, content_type=doc.content_type, meta=curr_meta, embedding=doc.embedding + ) new_document.embedding = doc.embedding if return_embedding is True else None new_document.embedding = doc.embedding if return_embedding is True else None @@ -705,7 +707,7 @@ class InMemoryDocumentStore(BaseDocumentStore): raise NotImplementedError("InMemoryDocumentStore does not support headers.") logger.warning( - """DEPRECATION WARNINGS: + """DEPRECATION WARNINGS: 1. delete_all_documents() method is deprecated, please use delete_documents method For more details, please refer to the issue: https://github.com/deepset-ai/haystack/issues/1045 """ diff --git a/haystack/document_stores/pinecone.py b/haystack/document_stores/pinecone.py index 0883db8886..cfb14a4763 100644 --- a/haystack/document_stores/pinecone.py +++ b/haystack/document_stores/pinecone.py @@ -392,7 +392,10 @@ class PineconeDocumentStore(BaseDocumentStore): elif duplicate_documents == "fail": # Otherwise, we raise an error raise DuplicateDocumentError(f"Duplicate document IDs found in batch: {ids}") - metadata = [{"content": doc.content, **doc.meta} for doc in document_objects[i : i + batch_size]] + metadata = [ + {"content": doc.content, "content_type": doc.content_type, **doc.meta} + for doc in document_objects[i : i + batch_size] + ] if add_vectors: embeddings = [doc.embedding for doc in document_objects[i : i + batch_size]] embeddings_to_index = np.array(embeddings, dtype="float32") @@ -498,7 +501,7 @@ class PineconeDocumentStore(BaseDocumentStore): metadata = [] ids = [] for doc in document_batch: - metadata.append({"content": doc.content, **doc.meta}) + metadata.append({"content": doc.content, "content_type": doc.content_type, **doc.meta}) ids.append(doc.id) # Update existing vectors in pinecone index self.pinecone_indexes[index].upsert( @@ -892,7 +895,7 @@ class PineconeDocumentStore(BaseDocumentStore): doc = self.get_documents_by_id(ids=[id], index=index, return_embedding=True)[0] if doc.embedding is not None: - meta = {"content": doc.content, **meta} + meta = {"content": doc.content, "content_type": doc.content_type, **meta} self.pinecone_indexes[index].upsert(vectors=[(id, doc.embedding.tolist(), meta)], namespace=namespace) def delete_documents( @@ -1177,7 +1180,8 @@ class PineconeDocumentStore(BaseDocumentStore): documents = [] for _id, meta in zip(ids, metadata): content = meta.pop("content") - doc = Document(id=_id, content=content, meta=meta) + content_type = meta.pop("content_type") + doc = Document(id=_id, content=content, content_type=content_type, meta=meta) documents.append(doc) if return_embedding: if values is None: diff --git a/haystack/document_stores/sql.py b/haystack/document_stores/sql.py index 4f88c842aa..21092c73b8 100644 --- a/haystack/document_stores/sql.py +++ b/haystack/document_stores/sql.py @@ -628,7 +628,7 @@ class SQLDocumentStore(BaseDocumentStore): raise NotImplementedError("SQLDocumentStore does not support headers.") logger.warning( - """DEPRECATION WARNINGS: + """DEPRECATION WARNINGS: 1. delete_all_documents() method is deprecated, please use delete_documents method For more details, please refer to the issue: https://github.com/deepset-ai/haystack/issues/1045 """ diff --git a/haystack/json-schemas/haystack-pipeline-1.10.0rc0.schema.json b/haystack/json-schemas/haystack-pipeline-1.10.0rc0.schema.json index 4a8569a8c3..37cbcef93f 100644 --- a/haystack/json-schemas/haystack-pipeline-1.10.0rc0.schema.json +++ b/haystack/json-schemas/haystack-pipeline-1.10.0rc0.schema.json @@ -2572,6 +2572,165 @@ ], "type": "object" }, + "MultiModalRetrieverComponent": { + "additionalProperties": false, + "properties": { + "name": { + "description": "Custom name for the component. Helpful for visualization and debugging.", + "title": "Name", + "type": "string" + }, + "params": { + "additionalProperties": false, + "description": "Each parameter can reference other components defined in the same YAML file.", + "properties": { + "batch_size": { + "default": 16, + "title": "Batch Size", + "type": "integer" + }, + "devices": { + "anyOf": [ + { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "string" + } + ] + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Devices" + }, + "document_embedding_models": { + "additionalProperties": { + "anyOf": [ + { + "format": "path", + "type": "string" + }, + { + "type": "string" + } + ] + }, + "title": "Document Embedding Models", + "type": "object" + }, + "document_feature_extractors_params": { + "additionalProperties": { + "type": "object" + }, + "default": { + "text": { + "max_length": 256 + } + }, + "title": "Document Feature Extractors Params", + "type": "object" + }, + "document_store": { + "title": "Document Store", + "type": "string" + }, + "embed_meta_fields": { + "default": [ + "name" + ], + "items": { + "type": "string" + }, + "title": "Embed Meta Fields", + "type": "array" + }, + "progress_bar": { + "default": true, + "title": "Progress Bar", + "type": "boolean" + }, + "query_embedding_model": { + "anyOf": [ + { + "format": "path", + "type": "string" + }, + { + "type": "string" + } + ], + "title": "Query Embedding Model" + }, + "query_feature_extractor_params": { + "default": { + "max_length": 64 + }, + "title": "Query Feature Extractor Params", + "type": "object" + }, + "query_type": { + "default": "text", + "title": "Query Type", + "type": "string" + }, + "scale_score": { + "default": true, + "title": "Scale Score", + "type": "boolean" + }, + "similarity_function": { + "default": "dot_product", + "title": "Similarity Function", + "type": "string" + }, + "top_k": { + "default": 10, + "title": "Top K", + "type": "integer" + }, + "use_auth_token": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Use Auth Token" + } + }, + "required": [ + "document_store", + "query_embedding_model", + "document_embedding_models" + ], + "title": "Parameters", + "type": "object" + }, + "type": { + "const": "MultiModalRetriever", + "description": "Haystack Class name for the component.", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "name" + ], + "type": "object" + }, "MultihopEmbeddingRetrieverComponent": { "additionalProperties": false, "properties": { @@ -6241,6 +6400,9 @@ { "$ref": "#/definitions/MarkdownConverterComponent" }, + { + "$ref": "#/definitions/MultiModalRetrieverComponent" + }, { "$ref": "#/definitions/MultihopEmbeddingRetrieverComponent" }, diff --git a/haystack/json-schemas/haystack-pipeline-main.schema.json b/haystack/json-schemas/haystack-pipeline-main.schema.json index 50bedf486a..3f5f36eb51 100644 --- a/haystack/json-schemas/haystack-pipeline-main.schema.json +++ b/haystack/json-schemas/haystack-pipeline-main.schema.json @@ -2572,6 +2572,165 @@ ], "type": "object" }, + "MultiModalRetrieverComponent": { + "additionalProperties": false, + "properties": { + "name": { + "description": "Custom name for the component. Helpful for visualization and debugging.", + "title": "Name", + "type": "string" + }, + "params": { + "additionalProperties": false, + "description": "Each parameter can reference other components defined in the same YAML file.", + "properties": { + "batch_size": { + "default": 16, + "title": "Batch Size", + "type": "integer" + }, + "devices": { + "anyOf": [ + { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "string" + } + ] + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Devices" + }, + "document_embedding_models": { + "additionalProperties": { + "anyOf": [ + { + "format": "path", + "type": "string" + }, + { + "type": "string" + } + ] + }, + "title": "Document Embedding Models", + "type": "object" + }, + "document_feature_extractors_params": { + "additionalProperties": { + "type": "object" + }, + "default": { + "text": { + "max_length": 256 + } + }, + "title": "Document Feature Extractors Params", + "type": "object" + }, + "document_store": { + "title": "Document Store", + "type": "string" + }, + "embed_meta_fields": { + "default": [ + "name" + ], + "items": { + "type": "string" + }, + "title": "Embed Meta Fields", + "type": "array" + }, + "progress_bar": { + "default": true, + "title": "Progress Bar", + "type": "boolean" + }, + "query_embedding_model": { + "anyOf": [ + { + "format": "path", + "type": "string" + }, + { + "type": "string" + } + ], + "title": "Query Embedding Model" + }, + "query_feature_extractor_params": { + "default": { + "max_length": 64 + }, + "title": "Query Feature Extractor Params", + "type": "object" + }, + "query_type": { + "default": "text", + "title": "Query Type", + "type": "string" + }, + "scale_score": { + "default": true, + "title": "Scale Score", + "type": "boolean" + }, + "similarity_function": { + "default": "dot_product", + "title": "Similarity Function", + "type": "string" + }, + "top_k": { + "default": 10, + "title": "Top K", + "type": "integer" + }, + "use_auth_token": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Use Auth Token" + } + }, + "required": [ + "document_store", + "query_embedding_model", + "document_embedding_models" + ], + "title": "Parameters", + "type": "object" + }, + "type": { + "const": "MultiModalRetriever", + "description": "Haystack Class name for the component.", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "name" + ], + "type": "object" + }, "MultihopEmbeddingRetrieverComponent": { "additionalProperties": false, "properties": { @@ -6241,6 +6400,9 @@ { "$ref": "#/definitions/MarkdownConverterComponent" }, + { + "$ref": "#/definitions/MultiModalRetrieverComponent" + }, { "$ref": "#/definitions/MultihopEmbeddingRetrieverComponent" }, diff --git a/haystack/modeling/data_handler/dataloader.py b/haystack/modeling/data_handler/dataloader.py index 6ebe9673ee..94f273af97 100644 --- a/haystack/modeling/data_handler/dataloader.py +++ b/haystack/modeling/data_handler/dataloader.py @@ -1,28 +1,33 @@ +from typing import List + from math import ceil + import torch -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, Dataset, Sampler + +from haystack.errors import ModelingError class NamedDataLoader(DataLoader): - """ - A modified version of the PyTorch DataLoader that returns a dictionary where the key is - the name of the tensor and the value is the tensor itself. - """ - - def __init__(self, dataset, batch_size, sampler=None, tensor_names=None, num_workers=0, pin_memory=False): + def __init__( + self, + dataset: Dataset, + batch_size: int, + sampler: Sampler = None, + tensor_names: List[str] = None, + num_workers: int = 0, + pin_memory: bool = False, + ): """ + A modified version of the PyTorch DataLoader that returns a dictionary where the key is + the name of the tensor and the value is the tensor itself. + :param dataset: The dataset that will be wrapped by this NamedDataLoader - :type dataset: Dataset :param sampler: The sampler used by the NamedDataLoader to choose which samples to include in the batch - :type sampler: Sampler :param batch_size: The size of the batch to be returned by the NamedDataLoader - :type batch_size: int :param tensor_names: The names of the tensor, in the order that the dataset returns them in. - :type tensor_names: list :param num_workers: number of workers to use for the DataLoader - :type num_workers: int :param pin_memory: argument for Data Loader to use page-locked memory for faster transfer of data to GPU - :type pin_memory: bool """ def collate_fn(batch): @@ -38,11 +43,10 @@ class NamedDataLoader(DataLoader): if type(batch[0]) == list: batch = batch[0] - assert len(batch[0]) == len( - _tensor_names - ), "Dataset contains {} tensors while there are {} tensor names supplied: {}".format( - len(batch[0]), len(_tensor_names), _tensor_names - ) + if len(batch[0]) != len(_tensor_names): + raise ModelingError( + f"Dataset contains {len(batch[0])} tensors while there are {len(_tensor_names)} tensor names supplied: {_tensor_names}" + ) lists_temp = [[] for _ in range(len(_tensor_names))] ret = dict(zip(_tensor_names, lists_temp)) @@ -55,7 +59,7 @@ class NamedDataLoader(DataLoader): return ret - super(NamedDataLoader, self).__init__( + super().__init__( dataset=dataset, sampler=sampler, batch_size=batch_size, @@ -71,19 +75,3 @@ class NamedDataLoader(DataLoader): return num_batches else: return super().__len__() - - -def covert_dataset_to_dataloader(dataset, sampler, batch_size): - """ - Wraps a PyTorch Dataset with a DataLoader. - - :param dataset: Dataset to be wrapped. - :type dataset: Dataset - :param sampler: PyTorch sampler used to pick samples in a batch. - :type sampler: Sampler - :param batch_size: Number of samples in the batch. - :return: A DataLoader that wraps the input Dataset. - """ - sampler_initialized = sampler(dataset) - data_loader = DataLoader(dataset, sampler=sampler_initialized, batch_size=batch_size) - return data_loader diff --git a/haystack/modeling/data_handler/processor.py b/haystack/modeling/data_handler/processor.py index 60117629a8..0a67fdc510 100644 --- a/haystack/modeling/data_handler/processor.py +++ b/haystack/modeling/data_handler/processor.py @@ -19,10 +19,9 @@ import requests from tqdm import tqdm from torch.utils.data import TensorDataset import transformers -from transformers import PreTrainedTokenizer +from transformers import PreTrainedTokenizer, AutoTokenizer -from haystack.modeling.model.tokenization import ( - get_tokenizer, +from haystack.modeling.model.feature_extraction import ( tokenize_batch_question_answering, tokenize_with_metadata, truncate_sequences, @@ -180,9 +179,11 @@ class Processor(ABC): "Loading tokenizer from deprecated config. " "If you used `custom_vocab` or `never_split_chars`, this won't work anymore." ) - tokenizer = get_tokenizer(load_dir, tokenizer_class=config["tokenizer"], do_lower_case=config["lower_case"]) + tokenizer = AutoTokenizer.from_pretrained( + load_dir, tokenizer_class=config["tokenizer"], do_lower_case=config["lower_case"] + ) else: - tokenizer = get_tokenizer(load_dir, tokenizer_class=config["tokenizer"]) + tokenizer = AutoTokenizer.from_pretrained(load_dir, tokenizer_class=config["tokenizer"]) # we have to delete the tokenizer string from config, because we pass it as Object del config["tokenizer"] @@ -218,7 +219,7 @@ class Processor(ABC): **kwargs, ): tokenizer_args = tokenizer_args or {} - tokenizer = get_tokenizer( + tokenizer = AutoTokenizer.from_pretrained( tokenizer_name_or_path, tokenizer_class=tokenizer_class, use_fast=use_fast, @@ -337,17 +338,7 @@ class Processor(ABC): :return: True if all the samples in the basket has computed its features, False otherwise """ - if basket.samples is None: - return False - elif len(basket.samples) == 0: - return False - if basket.samples is None: - return False - else: - for sample in basket.samples: - if sample.features is None: - return False - return True + return basket.samples and not any(sample.features is None for sample in basket.samples) def _log_samples(self, n_samples: int, baskets: List[SampleBasket]): logger.debug("*** Show %s random examples ***", n_samples) @@ -1350,9 +1341,15 @@ class TableTextSimilarityProcessor(Processor): processor_config_file = Path(load_dir) / "processor_config.json" config = json.load(open(processor_config_file)) # init tokenizer - query_tokenizer = get_tokenizer(load_dir, tokenizer_class=config["query_tokenizer"], subfolder="query") - passage_tokenizer = get_tokenizer(load_dir, tokenizer_class=config["passage_tokenizer"], subfolder="passage") - table_tokenizer = get_tokenizer(load_dir, tokenizer_class=config["table_tokenizer"], subfolder="table") + query_tokenizer = AutoTokenizer.from_pretrained( + load_dir, tokenizer_class=config["query_tokenizer"], subfolder="query" + ) + passage_tokenizer = AutoTokenizer.from_pretrained( + load_dir, tokenizer_class=config["passage_tokenizer"], subfolder="passage" + ) + table_tokenizer = AutoTokenizer.from_pretrained( + load_dir, tokenizer_class=config["table_tokenizer"], subfolder="table" + ) # we have to delete the tokenizer string from config, because we pass it as Object del config["query_tokenizer"] @@ -1978,7 +1975,7 @@ class InferenceProcessor(TextClassificationProcessor): processor_config_file = Path(load_dir) / "processor_config.json" config = json.load(open(processor_config_file)) # init tokenizer - tokenizer = get_tokenizer(load_dir, tokenizer_class=config["tokenizer"]) + tokenizer = AutoTokenizer.from_pretrained(load_dir, tokenizer_class=config["tokenizer"]) # we have to delete the tokenizer string from config, because we pass it as Object del config["tokenizer"] @@ -1999,37 +1996,6 @@ class InferenceProcessor(TextClassificationProcessor): ret: Dict = {} return ret - def dataset_from_dicts( - self, dicts: List[Dict], indices: List[int] = [], return_baskets: bool = False, debug: bool = False - ): - """ - Function to convert input dictionaries containing text into a torch dataset. - For normal operation with Language Models it calls the superclass' TextClassification.dataset_from_dicts method. - For slow tokenizers, s3e or wordembedding tokenizers the function works on _dict_to_samples and _sample_to_features - """ - # TODO remove this sections once tokenizers work the same way for slow/fast and our special tokenizers - if not self.tokenizer.is_fast: - self.baskets = [] - for d in dicts: - sample = self._dict_to_samples(dictionary=d) - features = self._sample_to_features(sample) - sample.features = features - basket = SampleBasket(id_internal=None, raw=d, id_external=None, samples=[sample]) - self.baskets.append(basket) - if indices and 0 not in indices: - pass - else: - self._log_samples(n_samples=1, baskets=self.baskets) - - problematic_ids: set = set() - dataset, tensornames = self._create_dataset() - ret = [dataset, tensornames, problematic_ids] - if return_baskets: - ret.append(self.baskets) - return ret - else: - return super().dataset_from_dicts(dicts=dicts, indices=indices, return_baskets=return_baskets, debug=debug) - # Private method to keep s3e pooling and embedding extraction working def _dict_to_samples(self, dictionary: Dict, **kwargs) -> Sample: # this tokenization also stores offsets diff --git a/haystack/modeling/model/tokenization.py b/haystack/modeling/model/feature_extraction.py similarity index 70% rename from haystack/modeling/model/tokenization.py rename to haystack/modeling/model/feature_extraction.py index 7ad2afc13a..979fd73d45 100644 --- a/haystack/modeling/model/tokenization.py +++ b/haystack/modeling/model/feature_extraction.py @@ -16,9 +16,18 @@ from typing import Dict, Any, Union, Tuple, Optional, List import re +import os +import json import logging +from pathlib import Path + import numpy as np -from transformers import AutoTokenizer, PreTrainedTokenizer, RobertaTokenizer +import transformers +from transformers import PreTrainedTokenizer, RobertaTokenizer, AutoConfig, AutoFeatureExtractor, AutoTokenizer + +# NOTE: These two constants are internals of HF. Keep in mind that they might be renamed or removed at any time. +from transformers.models.auto.feature_extraction_auto import FEATURE_EXTRACTOR_MAPPING_NAMES +from transformers.models.auto.tokenization_auto import TOKENIZER_MAPPING_NAMES from haystack.errors import ModelingError from haystack.modeling.data_handler.samples import SampleBasket @@ -31,45 +40,103 @@ logger = logging.getLogger(__name__) SPECIAL_TOKENIZER_CHARS = r"^(##|Ġ|▁)" -def get_tokenizer( - pretrained_model_name_or_path: str, - revision: str = None, - use_fast: bool = True, - use_auth_token: Optional[Union[str, bool]] = None, - **kwargs, -) -> PreTrainedTokenizer: - """ - Enables loading of different Tokenizer classes with a uniform interface. - Right now it always returns an instance of `AutoTokenizer`. +FEATURE_EXTRACTORS = { + **{key: AutoTokenizer for key in TOKENIZER_MAPPING_NAMES.keys()}, + **{key: AutoFeatureExtractor for key in FEATURE_EXTRACTOR_MAPPING_NAMES.keys()}, +} - :param pretrained_model_name_or_path: The path of the saved pretrained model or its name (e.g. `bert-base-uncased`) - :param revision: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash. - :param use_fast: Indicate if Haystack should try to load the fast version of the tokenizer (True) or use the Python one (False). Defaults to True. - :param use_auth_token: The API token used to download private models from Huggingface. - If this parameter is set to `True`, then the token generated when running - `transformers-cli login` (stored in ~/.huggingface) will be used. - Additional information can be found here - https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel.from_pretrained - :param kwargs: other kwargs to pass on to `PretrainedTokenizer.from_pretrained()` - :return: AutoTokenizer instance - """ - model_name_or_path = str(pretrained_model_name_or_path) - if "mlm" in model_name_or_path.lower(): - logging.error("MLM part of codebert is currently not supported in Haystack. Proceed at your own risk.") +DEFAULT_EXTRACTION_PARAMS = { + AutoTokenizer: { + "max_length": 256, + "add_special_tokens": True, + "truncation": True, + "truncation_strategy": "longest_first", + "padding": "max_length", + "return_token_type_ids": True, + }, + AutoFeatureExtractor: {"return_tensors": "pt"}, +} - params = {} - if any(tokenizer_type in model_name_or_path for tokenizer_type in ["albert", "xlnet"]): - params["keep_accents"] = True - return AutoTokenizer.from_pretrained( - pretrained_model_name_or_path=model_name_or_path, - revision=revision, - use_fast=use_fast, - use_auth_token=use_auth_token, - **params, +class FeatureExtractor: + def __init__( + self, + pretrained_model_name_or_path: Union[str, Path], + revision: str = None, + use_fast: bool = True, + use_auth_token: Optional[Union[str, bool]] = None, **kwargs, - ) + ): + """ + Enables loading of different feature extractors, including tokenizers, with a uniform interface. + + Use `FeatureExtractor.extract_features()` to convert your input queries, documents, images, and tables + into vectors that you can pass to the language model. + + :param pretrained_model_name_or_path: The path of the saved pretrained model or its name (for example, `bert-base-uncased`) + :param revision: The version of the model to use from the Hugging Face model hub. It can be tag name, branch name, or commit hash. + :param use_fast: Indicate if Haystack should try to load the fast version of the tokenizer (True) or use the Python one (False). Defaults to True. + :param use_auth_token: The API token used to download private models from Hugging Face. + If this parameter is set to `True`, then the token generated when running + `transformers-cli login` (stored in ~/.huggingface) is used. + For more information, see + [Hugging Face documentation](https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel.from_pretrained) + :param kwargs: Other kwargs you want to pass on to `PretrainedTokenizer.from_pretrained()` + """ + model_name_or_path = str(pretrained_model_name_or_path) + model_type = None + + config_file = Path(pretrained_model_name_or_path) / "tokenizer_config.json" + if os.path.exists(config_file): + # it's a local directory + config = json.load(open(config_file)) + feature_extractor_classname = config["tokenizer_class"] + logger.debug(f"⛏️ Selected feature extractor: {feature_extractor_classname} (from {config_file})") + # Use FastTokenizers as much as possible + try: + feature_extractor_class = getattr(transformers, feature_extractor_classname + "Fast") + logger.debug( + "Fast version of this tokenizer exists. Loaded class: %s", + feature_extractor_class.__class__.__name__, + ) + except AttributeError: + logger.debug("Fast version could not be loaded. Falling back to base version.") + feature_extractor_class = getattr(transformers, feature_extractor_classname) + + else: + # it's a HF Hub identifier + config = AutoConfig.from_pretrained( + pretrained_model_name_or_path=model_name_or_path, use_auth_token=use_auth_token, revision=revision + ) + model_type = config.model_type + try: + feature_extractor_class = FEATURE_EXTRACTORS[model_type] + except KeyError as e: + raise ModelingError( + f"'{pretrained_model_name_or_path}' has no known feature extractor. " + "Haystack can assign tokenizers to the following model types: " + # Using chr(10) instead of \n due to f-string limitation + # https://peps.python.org/pep-0498/#specification: "Backslashes may not appear anywhere within expressions" + f"\n- {f'{chr(10)}- '.join(FEATURE_EXTRACTORS.keys())}" + ) from e + logger.debug( + f"⛏️ Selected feature extractor: {feature_extractor_class.__name__} (for model type '{model_type}')" + ) + + self.default_params = DEFAULT_EXTRACTION_PARAMS.get(feature_extractor_class, {}) + + self.feature_extractor = feature_extractor_class.from_pretrained( + pretrained_model_name_or_path=model_name_or_path, + revision=revision, + use_fast=use_fast, + use_auth_token=use_auth_token, + **kwargs, + ) + + def __call__(self, **kwargs): + params = {**self.default_params, **(kwargs or {})} + return self.feature_extractor(**params) def tokenize_batch_question_answering( @@ -101,7 +168,11 @@ def tokenize_batch_question_answering( # # Tokenize texts in batch mode texts = [d["context"] for d in pre_baskets] tokenized_docs_batch = tokenizer( - texts, return_offsets_mapping=True, return_special_tokens_mask=True, add_special_tokens=False, verbose=False + text=texts, + return_offsets_mapping=True, + return_special_tokens_mask=True, + add_special_tokens=False, + verbose=False, ) # Extract relevant data diff --git a/haystack/modeling/model/language_model.py b/haystack/modeling/model/language_model.py index 020cec6731..f1a3a48e25 100644 --- a/haystack/modeling/model/language_model.py +++ b/haystack/modeling/model/language_model.py @@ -13,9 +13,9 @@ # 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. -""" +""" Acknowledgements: Many of the modeling parts here come from the great transformers repository: https://github.com/huggingface/transformers. -Thanks for the great work! +Thanks for the great work! """ from typing import Type, Optional, Dict, Any, Union, List @@ -26,7 +26,6 @@ import logging import os from abc import ABC, abstractmethod from pathlib import Path -from functools import wraps import numpy as np import torch from torch import nn @@ -36,6 +35,7 @@ from transformers import AutoModel, AutoConfig from transformers.modeling_utils import SequenceSummary from haystack.errors import ModelingError +from haystack.modeling.utils import silence_transformers_logs logger = logging.getLogger(__name__) @@ -58,30 +58,6 @@ LANGUAGE_HINTS = ( OUTPUT_DIM_NAMES = ["dim", "hidden_size", "d_model"] -def silence_transformers_logs(from_pretrained_func): - """ - A wrapper that raises the log level of Transformers to - ERROR to hide some unnecessary warnings. - """ - - @wraps(from_pretrained_func) - def quiet_from_pretrained_func(cls, *args, **kwargs): - - # Raise the log level of Transformers - t_logger = logging.getLogger("transformers") - original_log_level = t_logger.level - t_logger.setLevel(logging.ERROR) - - result = from_pretrained_func(cls, *args, **kwargs) - - # Restore the log level - t_logger.setLevel(original_log_level) - - return result - - return quiet_from_pretrained_func - - # TODO analyse if LMs can be completely used through HF transformers class LanguageModel(nn.Module, ABC): """ diff --git a/haystack/modeling/model/multimodal/__init__.py b/haystack/modeling/model/multimodal/__init__.py new file mode 100644 index 0000000000..63631ea257 --- /dev/null +++ b/haystack/modeling/model/multimodal/__init__.py @@ -0,0 +1,178 @@ +from typing import Optional, Union, Dict, Any, Type, List + +import logging +from pathlib import Path +from transformers import AutoConfig +import torch +from huggingface_hub import hf_hub_download + +from haystack.modeling.model.multimodal.base import HaystackModel +from haystack.modeling.model.multimodal.sentence_transformers import HaystackSentenceTransformerModel + + +logger = logging.getLogger(__name__) + + +#: Match the name of the HuggingFace Model class to the corresponding Haystack wrapper +HUGGINGFACE_TO_HAYSTACK: Dict[str, Type[HaystackModel]] = { + "CLIP": HaystackSentenceTransformerModel, + "MPNet": HaystackSentenceTransformerModel, + "DPRContextEncoder": HaystackSentenceTransformerModel, + "DPRQuestionEncoder": HaystackSentenceTransformerModel, +} + + +#: HF Capitalization pairs. Contains alternative capitalizations. +HUGGINGFACE_CAPITALIZE = { + "big-bird": "BigBird", + "deberta-v2": "DebertaV2", + "xlm-roberta": "XLMRoberta", + **{k.lower(): k for k in HUGGINGFACE_TO_HAYSTACK.keys()}, +} + + +def get_model( + pretrained_model_name_or_path: Union[Path, str], + content_type: str, # change to ContentTypes starting Python3.8 + devices: Optional[List[torch.device]] = None, + autoconfig_kwargs: Optional[Dict[str, Any]] = None, + model_kwargs: Optional[Dict[str, Any]] = None, + feature_extractor_kwargs: Optional[Dict[str, Any]] = None, + pooler_kwargs: Optional[Dict[str, Any]] = None, +) -> HaystackModel: + """ + Load a pretrained language model by specifying its name and either downloading the model from the Hugging Face hub + (if it's given a model identifier from Hugging Face Hub) or loading it from disk (if it's given a local path). + + For all supported model variations, see [Models](https://huggingface.co/models). + The appropriate language model class is inferred automatically from the model's configuration or its name. + + :param pretrained_model_name_or_path: The path of the saved pretrained model or its name. + :param content_type: The type (text, image, ...) of content the model should handle. + :param autoconfig_kwargs: Additional keyword arguments to pass to AutoConfig, like the revision or the auth key. + :param model_kwargs: Additional keyword arguments to pass to the language model constructor. + Haystack applies some default parameters to some models. You can override them by specifying the + desired value in this parameter. See `DEFAULT_MODEL_PARAMS`. + :param feature_extractor_kwargs: A dictionary of parameters to pass to the feature extractor's initialization (revision, use_auth_key, etc...) + Haystack applies some default parameters to some models. You can override them by specifying the + desired value in this parameter. See `DEFAULT_MODEL_PARAMS`. + :param pooler_kwargs: A dictionary of parameters to pass to the pooler's initialization (summary_last_dropout, summary_activation, etc...) + Haystack applies some default parameters to some models. You can override them by specifying the + desired value in this parameter. See `POOLER_PARAMETERS`. + """ + autoconfig_kwargs = autoconfig_kwargs or {} + model_kwargs = model_kwargs or {} + feature_extractor_kwargs = feature_extractor_kwargs or {} + pooler_kwargs = pooler_kwargs or {} + + if not pretrained_model_name_or_path or not isinstance(pretrained_model_name_or_path, (str, Path)): + raise ValueError( + f"{pretrained_model_name_or_path} is not a valid 'pretrained_model_name_or_path' value. " + "Please provide a string or a Path object." + ) + model_name = str(pretrained_model_name_or_path) + model_type: Optional[str] = "" + model_wrapper_class: Type[HaystackModel] + + # Prepare the kwargs the model wrapper expects (see each wrapper's init for details) + wrapper_kwarg_groups = {} + wrapper_kwarg_groups["model_kwargs"] = model_kwargs + + # SentenceTransformers are much faster, so use them whenever possible + if _is_sentence_transformers_model( + pretrained_model_name_or_path, use_auth_token=autoconfig_kwargs.get("use_auth_token", False) + ): + model_wrapper_class = HaystackSentenceTransformerModel + try: + # Use AutoConfig to log some more info about the model class + config = AutoConfig.from_pretrained(pretrained_model_name_or_path=model_name, **autoconfig_kwargs) + model_type = config.model_type + except Exception as e: + logger.debug(f"Can't find model type for {pretrained_model_name_or_path}: {e}") + + if feature_extractor_kwargs is not None: + logger.debug( + "Can't forward feature_extractor_kwargs to a SentenceTransformers model. " + "These kwargs are being dropped. " + f"Content of feature_extractor_kwargs: {feature_extractor_kwargs}" + ) + + else: + # Use AutoConfig to understand the model class + config = AutoConfig.from_pretrained(pretrained_model_name_or_path=model_name, **autoconfig_kwargs) + if not config.model_type: + logger.error( + f"Model type not understood for '{pretrained_model_name_or_path}'. Please provide the name of " + "a model that can be downloaded from the Model Hub.\nUsing the AutoModel class. " + "THIS CAN CAUSE CRASHES and won't work for models that are not working with text." + ) + model_type = None + else: + try: + model_type = HUGGINGFACE_CAPITALIZE[config.model_type.lower()] + except KeyError as e: + logger.error( + f"Haystack doesn't support model '{pretrained_model_name_or_path}' (type '{config.model_type.lower()}') " + "We'll use the AutoModel class for it. " + "THIS CAN CAUSE CRASHES and won't work for models that are not working with text. " + f"Supported model types: {', '.join(HUGGINGFACE_CAPITALIZE.keys())}" + ) + model_type = None + + # Find the HF class corresponding to this model type + try: + model_wrapper_class = HUGGINGFACE_TO_HAYSTACK[model_type or "AutoModel"] + except KeyError as e: + raise ValueError( + f"The type of the given model (name/path: {pretrained_model_name_or_path}, detected type: {model_type or config.model_type}) " + "is not supported by Haystack or was not correctly identified. Please use supported models only. " + f"Supported model types: {', '.join(HUGGINGFACE_TO_HAYSTACK.keys())}" + ) from e + + if feature_extractor_kwargs: + wrapper_kwarg_groups["feature_extractor_kwargs"] = feature_extractor_kwargs + if pooler_kwargs: + wrapper_kwarg_groups["pooler_kwargs"] = pooler_kwargs + + # Instantiate the model's wrapper + model_wrapper = model_wrapper_class( + pretrained_model_name_or_path=pretrained_model_name_or_path, + model_type=model_type, + content_type=content_type, + **wrapper_kwarg_groups, + ) + model_wrapper.to(devices) + + return model_wrapper + + +def _is_sentence_transformers_model(pretrained_model_name_or_path: Union[Path, str], use_auth_token: Union[bool, str]): + + # Check if sentence transformers config file is in local path + if Path(pretrained_model_name_or_path).exists(): + if (Path(pretrained_model_name_or_path) / "config_sentence_transformers.json").exists(): + return True + + # Check if sentence transformers config file is in model hub + try: + hf_hub_download( + repo_id=str(pretrained_model_name_or_path), + filename="config_sentence_transformers.json", + use_auth_token=use_auth_token, + ) + return True + + except Exception as e: + logger.debug("%s not found in model hub: an error occurred. Error: %s", pretrained_model_name_or_path, e) + + # Pattern matching the name as a last resort + if str(pretrained_model_name_or_path).startswith("sentence-transformers"): + logger.debug( + "The model name starts with 'sentence-transformers': assuming this is a Sentence Transformers model." + ) + return True + + logger.debug( + "The model name doesn't start with 'sentence-transformers': assuming this is NOT a Sentence Transformers model." + ) + return False diff --git a/haystack/modeling/model/multimodal/base.py b/haystack/modeling/model/multimodal/base.py new file mode 100644 index 0000000000..bfcea5b054 --- /dev/null +++ b/haystack/modeling/model/multimodal/base.py @@ -0,0 +1,56 @@ +from typing import Any, List, Union, Optional + +import logging +from pathlib import Path +from abc import ABC, abstractmethod + +import torch + + +logger = logging.getLogger(__name__) + + +class HaystackModel(ABC): + """ + The interface on top of HaystackTransformer and HaystackSentenceTransformer. + """ + + def __init__( + self, pretrained_model_name_or_path: Union[str, Path], model_type: Optional[str], content_type: str + ): # replace the type of content_type with ContentTypes starting Python3.8 + """ + :param pretrained_model_name_or_path: The name of the model to load + :param model_type: the value of `model_type` from the model's `Config` class. + :param content_type: The type of data (such as "text", "image" and so on) the model should process. + See the values of `haystack.schema.ContentTypes`. + """ + logger.info( + f" 🤖 Loading '{pretrained_model_name_or_path}' " + f"({self.__class__.__name__} of type '{model_type if model_type else ''}' " + f"for {content_type} data)" + ) + self.model_name_or_path = pretrained_model_name_or_path + self.model_type = model_type + self.content_type = content_type + + @abstractmethod + def encode(self, data: List[Any], **kwargs) -> torch.Tensor: + """ + Run the model on the input data to obtain output vectors. + """ + raise NotImplementedError("Abstract method, use a subclass.") + + @abstractmethod + def to(self, devices: Optional[List[torch.device]]) -> None: + """ + Send the model to the specified PyTorch device(s) + """ + raise NotImplementedError("Abstract method, use a subclass.") + + @property + @abstractmethod + def embedding_dim(self) -> int: + """ + The output embedding size. + """ + raise NotImplementedError("Abstract method, use a subclass.") diff --git a/haystack/modeling/model/multimodal/sentence_transformers.py b/haystack/modeling/model/multimodal/sentence_transformers.py new file mode 100644 index 0000000000..e0b5c47d40 --- /dev/null +++ b/haystack/modeling/model/multimodal/sentence_transformers.py @@ -0,0 +1,92 @@ +from typing import Optional, Dict, Any, List, Union + +import logging +from pathlib import Path + +import torch +from torch import nn +from sentence_transformers import SentenceTransformer + +from haystack.modeling.model.multimodal.base import HaystackModel +from haystack.schema import ContentTypes + + +logger = logging.getLogger(__name__) + + +class HaystackSentenceTransformerModel(HaystackModel): + """ + Parent class for `sentence-transformers` models. + + These models read raw data (text, image), internally extract features, and return vectors that capture the meaning + of the original data. + + Models inheriting from `HaystackSentenceTransformerModel` are designed to be used in parallel one with the other + in multimodal retrieval settings, for example image retrieval from a text query, mixed table/text retrieval, and so on. + """ + + def __init__( + self, + pretrained_model_name_or_path: Union[str, Path], + model_type: str, + content_type: ContentTypes, + model_kwargs: Optional[Dict[str, Any]] = None, + ): + """ + :param pretrained_model_name_or_path: The name of the model to load. + :param model_type: the value of `model_type` from the model's `Config` class + :param content_type: The type of data (such as "text", "image", and so on) the model should process. + See the values of `haystack.schema.ContentTypes`. + :param model_kwargs: A dictionary of parameters to pass to the model's initialization. + (revision, use_auth_key, and so on) + Haystack applies some default parameters to some models. You can override them by specifying the + desired value in this parameter. See `DEFAULT_MODEL_PARAMS`. + """ + super().__init__( + pretrained_model_name_or_path=pretrained_model_name_or_path, + model_type=model_type, + content_type=content_type, + ) + try: + self.model = SentenceTransformer(pretrained_model_name_or_path, **(model_kwargs or {})) + except Exception as e: + logger.exception( + "Models of type '%s' like %s " + "are only supported through sentence-transformers. Make sure this " + "model is compatible with sentence-transformers or use an alternative, compatible " + "implementation of this model.", + model_type, + pretrained_model_name_or_path, + ) + + @property + def embedding_dim(self) -> int: + """ + Finds out the output embedding dim by running the model on a minimal amount of mock data. + """ + emedding_dim = self.model.get_sentence_embedding_dimension() + if not emedding_dim: + logger.warning( + "Can't find the output embedding dimensions for '%s'. Some checks will not run as intended.", + self.model_name_or_path, + ) + return emedding_dim + + def to(self, devices: Optional[List[torch.device]]) -> None: + """ + Send the model to the specified PyTorch devices. + """ + if devices: + if len(devices) > 1: + self.model = nn.DataParallel(self.model, device_ids=devices) + else: + self.model.to(devices[0]) + + def encode(self, data: List[Any], **kwargs) -> torch.Tensor: + """ + Generate the tensors representing the input data. + + Validates the inputs according to what the subclass declared in the `expected_inputs` property. + Then passes the vectors to the `_forward()` method and returns its output untouched. + """ + return self.model.encode(data, convert_to_tensor=True, **kwargs) diff --git a/haystack/modeling/utils.py b/haystack/modeling/utils.py index 755082a052..14119172c9 100644 --- a/haystack/modeling/utils.py +++ b/haystack/modeling/utils.py @@ -5,8 +5,10 @@ import os import pickle import random import signal +from functools import wraps from copy import deepcopy from itertools import islice + import numpy as np import torch import torch.distributed as dist @@ -28,6 +30,30 @@ class GracefulKiller: self.kill_now = True +def silence_transformers_logs(from_pretrained_func): + """ + A wrapper that raises the log level of Transformers to + ERROR to hide some unnecessary warnings. + """ + + @wraps(from_pretrained_func) + def quiet_from_pretrained_func(cls, *args, **kwargs): + + # Raise the log level of Transformers + t_logger = logging.getLogger("transformers") + original_log_level = t_logger.level + t_logger.setLevel(logging.ERROR) + + result = from_pretrained_func(cls, *args, **kwargs) + + # Restore the log level + t_logger.setLevel(original_log_level) + + return result + + return quiet_from_pretrained_func + + def set_all_seeds(seed: int, deterministic_cudnn: bool = False) -> None: """ Setting multiple seeds to make runs reproducible. diff --git a/haystack/nodes/__init__.py b/haystack/nodes/__init__.py index 6fd9b63c41..4471b723f2 100644 --- a/haystack/nodes/__init__.py +++ b/haystack/nodes/__init__.py @@ -40,6 +40,7 @@ from haystack.nodes.retriever import ( TfidfRetriever, Text2SparqlRetriever, TableTextRetriever, + MultiModalRetriever, ) from haystack.nodes.summarizer import BaseSummarizer, TransformersSummarizer from haystack.nodes.translator import BaseTranslator, TransformersTranslator diff --git a/haystack/nodes/retriever/__init__.py b/haystack/nodes/retriever/__init__.py index c7c261a045..5e1c6f1f19 100644 --- a/haystack/nodes/retriever/__init__.py +++ b/haystack/nodes/retriever/__init__.py @@ -14,3 +14,4 @@ from haystack.nodes.retriever.sparse import ( TfidfRetriever, ) from haystack.nodes.retriever.text2sparql import Text2SparqlRetriever +from haystack.nodes.retriever.multimodal import MultiModalRetriever diff --git a/haystack/nodes/retriever/multimodal/__init__.py b/haystack/nodes/retriever/multimodal/__init__.py new file mode 100644 index 0000000000..1f445c8e94 --- /dev/null +++ b/haystack/nodes/retriever/multimodal/__init__.py @@ -0,0 +1,2 @@ +from haystack.nodes.retriever.multimodal.retriever import MultiModalRetriever +from haystack.nodes.retriever.multimodal.embedder import MultiModalEmbedder diff --git a/haystack/nodes/retriever/multimodal/embedder.py b/haystack/nodes/retriever/multimodal/embedder.py new file mode 100644 index 0000000000..7c3e6d0a7b --- /dev/null +++ b/haystack/nodes/retriever/multimodal/embedder.py @@ -0,0 +1,207 @@ +from typing import Union, Optional, Dict, List, Any + +import logging +from pathlib import Path + +import torch +from tqdm import tqdm +import numpy as np +from PIL import Image + +from haystack.modeling.model.multimodal import get_model +from haystack.errors import NodeError, ModelingError +from haystack.modeling.model.multimodal.base import HaystackModel +from haystack.schema import Document +from haystack.utils.torch_utils import get_devices + + +logger = logging.getLogger(__name__) + + +class MultiModalRetrieverError(NodeError): + pass + + +FilterType = Dict[str, Union[Dict[str, Any], List[Any], str, int, float, bool]] + + +# TODO the keys should match with ContentTypes (currently 'audio' is missing) +DOCUMENT_CONVERTERS = { + # NOTE: Keep this '?' cleaning step, it needs to be double-checked for impact on the inference results. + "text": lambda doc: doc.content[:-1] if doc.content[-1] == "?" else doc.content, + "table": lambda doc: " ".join( + doc.content.columns.tolist() + [cell for row in doc.content.values.tolist() for cell in row] + ), + "image": lambda doc: Image.open(doc.content), +} + +CAN_EMBED_META = ["text", "table"] + + +class MultiModalEmbedder: + def __init__( + self, + embedding_models: Dict[str, Union[Path, str]], # replace str with ContentTypes starting from Python3.8 + feature_extractors_params: Dict[str, Dict[str, Any]] = None, + batch_size: int = 16, + embed_meta_fields: List[str] = ["name"], + progress_bar: bool = True, + devices: Optional[List[Union[str, torch.device]]] = None, + use_auth_token: Optional[Union[str, bool]] = None, + ): + """ + Init the Retriever and all its models from a local or remote model checkpoint. + The checkpoint format matches the Hugging Face transformers' model format. + + :param embedding_models: A dictionary matching a local path or remote name of encoder checkpoint with + the content type it should handle ("text", "table", "image", etc...). + The format is the one that Hugging Face Hub models use. + Expected input format: `{'text': 'name_or_path_to_text_model', 'image': 'name_or_path_to_image_model', ... }` + Keep in mind that the models should output in the same embedding space for this retriever to work. + :param feature_extractors_params: A dictionary matching a content type ("text", "table", "image" and so on) with the + parameters of its own feature extractor if the model requires one. + Expected input format: `{'text': {'param_name': 'param_value', ...}, 'image': {'param_name': 'param_value', ...}, ...}` + :param batch_size: Number of questions or passages to encode at once. In case of multiple GPUs, this will be the total batch size. + :param embed_meta_fields: Concatenate the provided meta fields and text passage / image to a text pair that is + then used to create the embedding. + This is the approach used in the original paper and is likely to improve + performance if your titles contain meaningful information for retrieval + (topic, entities etc.). + :param progress_bar: Whether to show a tqdm progress bar or not. + Can be helpful to disable in production deployments to keep the logs clean. + :param devices: List of GPU (or CPU) devices to limit inference to certain GPUs and not use all available ones. + These strings are converted into pytorch devices, so use the string notation described + [in the pytorch documentation](https://pytorch.org/docs/simage/tensor_attributes.html?highlight=torch%20device#torch.torch.device) + (for example, ["cuda:0"]). Note: as multi-GPU training is currently not implemented for TableTextRetriever, + training only uses the first device provided in this list. + :param use_auth_token: API token used to download private models from Hugging Face. If this parameter is set to `True`, + the local token is used, which must be previously created using `transformer-cli login`. + For more information, see [Hugging Face documentation](https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel.from_pretrained) + """ + super().__init__() + + self.devices = get_devices(devices) + if batch_size < len(self.devices): + logger.warning("Batch size is lower than the number of devices. Not all GPUs will be utilized.") + + self.batch_size = batch_size + self.progress_bar = progress_bar + self.embed_meta_fields = embed_meta_fields + + feature_extractors_params = { + content_type: {"max_length": 256, **(feature_extractors_params or {}).get(content_type, {})} + for content_type in ["text", "table", "image", "audio"] # FIXME get_args(ContentTypes) from Python3.8 on + } + + self.models: Dict[str, HaystackModel] = {} # replace str with ContentTypes starting from Python3.8 + for content_type, embedding_model in embedding_models.items(): + self.models[content_type] = get_model( + pretrained_model_name_or_path=embedding_model, + content_type=content_type, + devices=self.devices, + autoconfig_kwargs={"use_auth_token": use_auth_token}, + model_kwargs={"use_auth_token": use_auth_token}, + feature_extractor_kwargs=feature_extractors_params[content_type], + ) + + # Check embedding sizes for models: they must all match + if len(self.models) > 1: + sizes = {model.embedding_dim for model in self.models.values()} + if None in sizes: + logger.warning( + "Haystack could not find the output embedding dimensions for '%s'. " + "Dimensions won't be checked before computing the embeddings.", + ", ".join( + {str(model.model_name_or_path) for model in self.models.values() if model.embedding_dim is None} + ), + ) + elif len(sizes) > 1: + embedding_sizes: Dict[int, List[str]] = {} + for model in self.models.values(): + embedding_sizes[model.embedding_dim] = embedding_sizes.get(model.embedding_dim, []) + [ + str(model.model_name_or_path) + ] + raise ValueError(f"Not all models have the same embedding size: {embedding_sizes}") + + def embed(self, documents: List[Document], batch_size: Optional[int] = None) -> np.ndarray: + """ + Create embeddings for a list of documents using the relevant encoder for their content type. + + :param documents: Documents to embed. + :return: Embeddings, one per document, in the form of a np.array + """ + batch_size = batch_size if batch_size is not None else self.batch_size + + all_embeddings = [] + for batch_index in tqdm( + iterable=range(0, len(documents), batch_size), + unit=" Docs", + desc=f"Create embeddings", + position=1, + leave=False, + disable=not self.progress_bar, + ): + docs_batch = documents[batch_index : batch_index + batch_size] + data_by_type = self._docs_to_data(documents=docs_batch) + + # Get output for each model + outputs_by_type: Dict[str, torch.Tensor] = {} # replace str with ContentTypes starting Python3.8 + for data_type, data in data_by_type.items(): + + model = self.models.get(data_type) + if not model: + raise ModelingError( + f"Some data of type {data_type} was passed, but no model capable of handling such data was " + f"initialized. Initialized models: {', '.join(self.models.keys())}" + ) + outputs_by_type[data_type] = model.encode(data=data) + + # Check the output sizes + embedding_sizes = [output.shape[-1] for output in outputs_by_type.values()] + + if not all(embedding_size == embedding_sizes[0] for embedding_size in embedding_sizes): + raise ModelingError( + "Some of the models are using a different embedding size. They should all match. " + f"Embedding sizes by model: " + f"{ {name: output.shape[-1] for name, output in outputs_by_type.items()} }" + ) + + # Combine the outputs in a single matrix + outputs = torch.stack(list(outputs_by_type.values())) + embeddings = outputs.view(-1, embedding_sizes[0]) + embeddings = embeddings.cpu() + all_embeddings.append(embeddings) + + return np.concatenate(all_embeddings) + + def _docs_to_data( + self, documents: List[Document] + ) -> Dict[str, List[Any]]: # FIXME replace str to ContentTypes from Python3.8 + """ + Extract the data to embed from each document and return them classified by content type. + + :param documents: The documents to prepare fur multimodal embedding. + :return: A dictionary containing one key for each content type, and a list of data extracted + from each document, ready to be passed to the feature extractor (for example the content + of a text document, a linearized table, a PIL image object, and so on) + """ + docs_data: Dict[str, List[Any]] = { # FIXME replace str to ContentTypes from Python3.8 + key: [] for key in ["text", "table", "image", "audio"] + } # FIXME get_args(ContentTypes) from Python3.8 on + for doc in documents: + try: + document_converter = DOCUMENT_CONVERTERS[doc.content_type] + except KeyError as e: + raise MultiModalRetrieverError( + f"Unknown content type '{doc.content_type}'. Known types: 'text', 'table', 'image'." # FIXME {', '.join(get_args(ContentTypes))}" from Python3.8 on + ) from e + + data = document_converter(doc) + + if doc.content_type in CAN_EMBED_META: + meta = [v for k, v in (doc.meta or {}).items() if k in self.embed_meta_fields] + data = f"{' '.join(meta)} {data}" if meta else data + + docs_data[doc.content_type].append(data) + + return {key: values for key, values in docs_data.items() if values} diff --git a/haystack/nodes/retriever/multimodal/retriever.py b/haystack/nodes/retriever/multimodal/retriever.py new file mode 100644 index 0000000000..0a27e667a7 --- /dev/null +++ b/haystack/nodes/retriever/multimodal/retriever.py @@ -0,0 +1,215 @@ +from typing import Union, Optional, Dict, List, Any + +import logging +from pathlib import Path + +import torch +import numpy as np + +from haystack.nodes.retriever import BaseRetriever +from haystack.document_stores import BaseDocumentStore +from haystack.schema import ContentTypes, Document +from haystack.nodes.retriever.multimodal.embedder import MultiModalEmbedder, MultiModalRetrieverError + + +logger = logging.getLogger(__name__) + + +FilterType = Optional[Dict[str, Union[Dict[str, Any], List[Any], str, int, float, bool]]] + + +class MultiModalRetriever(BaseRetriever): + def __init__( + self, + document_store: BaseDocumentStore, + query_embedding_model: Union[Path, str], + document_embedding_models: Dict[str, Union[Path, str]], # Replace str with ContentTypes starting Python3.8 + query_type: str = "text", # Replace str with ContentTypes starting Python3.8 + query_feature_extractor_params: Dict[str, Any] = {"max_length": 64}, + document_feature_extractors_params: Dict[str, Dict[str, Any]] = {"text": {"max_length": 256}}, + top_k: int = 10, + batch_size: int = 16, + embed_meta_fields: List[str] = ["name"], + similarity_function: str = "dot_product", + progress_bar: bool = True, + devices: Optional[List[Union[str, torch.device]]] = None, + use_auth_token: Optional[Union[str, bool]] = None, + scale_score: bool = True, + ): + """ + Retriever that uses a multiple encoder to jointly retrieve among a database consisting of different + data types. See the original paper for more details: + Kostić, Bogdan, et al. (2021): "Multi-modal Retrieval of Tables and Texts Using Tri-encoder Models" + (https://arxiv.org/abs/2108.04049), + + :param document_store: An instance of DocumentStore from which to retrieve documents. + :param query_embedding_model: Local path or remote name of question encoder checkpoint. The format equals the + one used by Hugging Face transformers' modelhub models. + :param document_embedding_models: Dictionary matching a local path or remote name of document encoder + checkpoint with the content type it should handle ("text", "table", "image", and so on). + The format equals the one used by Hugging Face transformers' modelhub models. + :param query_type: The content type of the query ("text", "image" and so on) + :param query_feature_extraction_params: The parameters to pass to the feature extractor of the query. + :param document_feature_extraction_params: The parameters to pass to the feature extractor of the documents. + :param top_k: How many documents to return per query. + :param batch_size: Number of questions or documents to encode at once. In case of multiple GPUs, this will be + the total batch size. + :param embed_meta_fields: Concatenate the provided meta fields to a (text) pair that is then used to create + the embedding. This is likely to improve performance if your titles contain meaningful information + for retrieval (topic, entities, and so on). Note that only text and table documents support this feature. + :param similarity_function: Which function to apply for calculating the similarity of query and document + embeddings during training. Options: `dot_product` (default) or `cosine`. + :param progress_bar: Whether to show a tqdm progress bar or not. + Can be helpful to disable in production deployments to keep the logs clean. + :param devices: List of GPU (or CPU) devices to limit inference to certain GPUs and not use all available ones. + These strings will be converted into pytorch devices, so use the string notation described in [Tensor Attributes] + (https://pytorch.org/docs/simage/tensor_attributes.html?highlight=torch%20device#torch.torch.device) + (e.g. ["cuda:0"]). Note: As multi-GPU training is currently not implemented for TableTextRetriever, + training only uses the first device provided in this list. + :param use_auth_token: API token used to download private models from Hugging Face. If this parameter is set + to `True`, the local token is used, which must be previously created using `transformer-cli login`. + For more information, see + [Hugging Face documentation](https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel.from_pretrained) + :param scale_score: Whether to scale the similarity score to the unit interval (range of [0,1]). + If true (default) similarity scores (e.g. cosine or dot_product) which naturally have a different value + range are scaled to a range of [0,1], where 1 means extremely relevant. + Otherwise raw similarity scores (for example, cosine or dot_product) are used. + """ + super().__init__() + + self.similarity_function = similarity_function + self.progress_bar = progress_bar + self.top_k = top_k + self.scale_score = scale_score + + self.document_embedder = MultiModalEmbedder( + embedding_models=document_embedding_models, + feature_extractors_params=document_feature_extractors_params, + batch_size=batch_size, + embed_meta_fields=embed_meta_fields, + progress_bar=progress_bar, + devices=devices, + use_auth_token=use_auth_token, + ) + + # Try to reuse the same embedder for queries if there is overlap + if document_embedding_models.get(query_type, None) == query_embedding_model: + self.query_embedder = self.document_embedder + else: + self.query_embedder = MultiModalEmbedder( + embedding_models={query_type: query_embedding_model}, + feature_extractors_params={query_type: query_feature_extractor_params}, + batch_size=batch_size, + embed_meta_fields=embed_meta_fields, + progress_bar=progress_bar, + devices=devices, + use_auth_token=use_auth_token, + ) + + self.document_store = document_store + + def retrieve( # type: ignore + self, + query: Any, + query_type: ContentTypes = "text", + filters: Optional[FilterType] = None, + top_k: Optional[int] = None, + index: str = None, + headers: Optional[Dict[str, str]] = None, + scale_score: bool = None, + ) -> List[Document]: + """ + Scan through documents in DocumentStore and return a small number of documents that are most relevant to the + supplied query. Returns a list of Documents. + + :param query: Query value. It might be text, a path, a table, and so on. + :param query_type: Type of the query ("text", "table", "image" and so on). + :param filters: Optional filters to narrow down the search space to documents whose metadata fulfill certain + conditions. It can be a single filter applied to each query or a list of filters + (one filter per query). + :param top_k: How many documents to return per query. Must be > 0. + :param index: The name of the index in the DocumentStore from which to retrieve documents. + :param batch_size: Number of queries to embed at a time. Must be > 0. + :param scale_score: Whether to scale the similarity score to the unit interval (range of [0,1]). + If true, similarity scores (for example, cosine or dot_product) which naturally have a different + value range is scaled to a range of [0,1], where 1 means extremely relevant. + Otherwise raw similarity scores (for example, cosine or dot_product) are used. + """ + return self.retrieve_batch( + queries=[query], + queries_type=query_type, + filters=[filters], + top_k=top_k, + index=index, + headers=headers, + batch_size=1, + scale_score=scale_score, + )[0] + + def retrieve_batch( # type: ignore + self, + queries: List[Any], + queries_type: ContentTypes = "text", + filters: Union[None, FilterType, List[FilterType]] = None, + top_k: Optional[int] = None, + index: str = None, + headers: Optional[Dict[str, str]] = None, + batch_size: Optional[int] = None, + scale_score: bool = None, + ) -> List[List[Document]]: + """ + Scan through documents in DocumentStore and return a small number of documents that are most relevant to the + supplied queries. Returns a list of lists of Documents (one list per query). + + This method assumes all queries are of the same data type. Mixed-type query batches (for example one image and one text) + are currently not supported. Group the queries by type and call `retrieve()` on uniform batches only. + + :param queries: List of query values. They might be text, paths, tables, and so on. + :param queries_type: Type of the query ("text", "table", "image" and so on) + :param filters: Optional filters to narrow down the search space to documents whose metadata fulfill certain + conditions. It can be a single filter that will be applied to each query or a list of filters + (one filter per query). + :param top_k: How many documents to return per query. Must be > 0. + :param index: The name of the index in the DocumentStore from which to retrieve documents. + :param batch_size: Number of queries to embed at a time. Must be > 0. + :param scale_score: Whether to scale the similarity score to the unit interval (range of [0,1]). + If True, similarity scores (for example, cosine or dot_product) which naturally have a different + value range are scaled to a range of [0,1], where 1 means extremely relevant. + Otherwise raw similarity scores (for example, cosine or dot_product) are used. + """ + filters_list: List[FilterType] + if not isinstance(filters, list): + filters_list = [filters] * len(queries) + else: + if len(filters) != len(queries): + raise MultiModalRetrieverError( + "The number of filters does not match the number of queries. Provide as many filters " + "as queries, or a single filter that will be applied to all queries." + ) + filters_list = filters + + top_k = top_k or self.top_k + index = index or self.document_store.index + scale_score = scale_score or self.scale_score + + # Embed the queries - we need them into Document format to leverage MultiModalEmbedder.embed() + query_docs = [Document(content=query, content_type=queries_type) for query in queries] + query_embeddings = self.query_embedder.embed(documents=query_docs, batch_size=batch_size) + + # Query documents by embedding (the actual retrieval step) + documents = [] + for query_embedding, query_filters in zip(query_embeddings, filters_list): + docs = self.document_store.query_by_embedding( + query_emb=query_embedding, + top_k=top_k, + filters=query_filters, + index=index, + headers=headers, + scale_score=scale_score, + ) + + documents.append(docs) + return documents + + def embed_documents(self, docs: List[Document]) -> np.ndarray: + return self.document_embedder.embed(documents=docs) diff --git a/haystack/schema.py b/haystack/schema.py index 1138c09b91..2e2d2d0964 100644 --- a/haystack/schema.py +++ b/haystack/schema.py @@ -34,12 +34,15 @@ logger = logging.getLogger(__name__) BaseConfig.arbitrary_types_allowed = True +#: Types of content_types supported +ContentTypes = Literal["text", "table", "image", "audio"] + @dataclass class Document: id: str content: Union[str, pd.DataFrame] - content_type: Literal["text", "table", "image", "audio"] = Field(default="text") + content_type: ContentTypes = Field(default="text") meta: Dict[str, Any] = Field(default={}) score: Optional[float] = None embedding: Optional[np.ndarray] = None @@ -52,7 +55,7 @@ class Document: def __init__( self, content: Union[str, pd.DataFrame], - content_type: Literal["text", "table", "image", "audio"] = "text", + content_type: ContentTypes = "text", id: Optional[str] = None, score: Optional[float] = None, meta: Optional[Dict[str, Any]] = None, @@ -69,7 +72,7 @@ class Document: It's particularly helpful for handling of duplicates and referencing documents in other objects (e.g. Labels) There's an easy option to convert from/to dicts via `from_dict()` and `to_dict`. :param content: Content of the document. For most cases, this will be text, but it can be a table or image. - :param content_type: One of "text", "table" or "image". Haystack components can use this to adjust their + :param content_type: One of "text", "table", "image" or "audio". Haystack components can use this to adjust their handling of Documents and check compatibility. :param id: Unique ID for the document. If not supplied by the user, we'll generate one automatically by creating a hash from the supplied text. This behaviour can be further adjusted by `id_hash_keys`. @@ -226,8 +229,8 @@ class Document: def __repr__(self): values = self.to_dict() - if values.get("embedding", False): - values["embedding"] = f"" + if values.get("embedding", None) is not None: + values["embedding"] = f"" return f"" def __str__(self): diff --git a/haystack/utils/torch_utils.py b/haystack/utils/torch_utils.py index c0e944fba0..a5aa6195b7 100644 --- a/haystack/utils/torch_utils.py +++ b/haystack/utils/torch_utils.py @@ -1,3 +1,6 @@ +from typing import Optional, List, Union + +import torch from torch.utils.data import Dataset @@ -10,3 +13,15 @@ class ListDataset(Dataset): def __getitem__(self, i): return self.original_list[i] + + +def get_devices(devices: Optional[List[Union[str, torch.device]]]) -> List[torch.device]: + """ + Convert a list of device names into a list of Torch devices, + depending on the system's configuration and hardware. + """ + if devices is not None: + return [torch.device(device) for device in devices] + elif torch.cuda.is_available(): + return [torch.device(device) for device in range(torch.cuda.device_count())] + return [torch.device("cpu")] diff --git a/pyproject.toml b/pyproject.toml index a59ba3f151..ed5c72050b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -327,9 +327,9 @@ min-similarity-lines=6 minversion = "6.0" addopts = "--strict-markers" markers = [ - "integration: integration tests", "unit: unit tests", - + "integration: integration tests", + "generator: generator tests", "summarizer: summarizer tests", "embedding_dim: uses a document store with non-default embedding dimension (e.g @pytest.mark.embedding_dim(128))", diff --git a/test/modeling/test_tokenization.py b/test/modeling/_test_feature_extraction_end2end.py similarity index 88% rename from test/modeling/test_tokenization.py rename to test/modeling/_test_feature_extraction_end2end.py index e755b77995..83651621f8 100644 --- a/test/modeling/test_tokenization.py +++ b/test/modeling/_test_feature_extraction_end2end.py @@ -9,7 +9,7 @@ from unittest.mock import MagicMock from tokenizers.pre_tokenizers import WhitespaceSplit import haystack -from haystack.modeling.model.tokenization import get_tokenizer +from haystack.modeling.model.feature_extraction import FeatureExtractor BERT = "bert-base-cased" @@ -28,8 +28,8 @@ SENTENCE_WITH_ELLIPSIS = "This is a sentence..." SENTENCE_WITH_LINEBREAK_1 = "and another one\n\n\nwithout space" SENTENCE_WITH_LINEBREAK_2 = """This is a sentence. With linebreak""" -SENTENCE_WITH_LINEBREAKS = """Sentence - with +SENTENCE_WITH_LINEBREAKS = """Sentence + with multiple newlines """ @@ -75,21 +75,21 @@ def convert_offset_from_word_reference_to_text_reference(offsets, words, word_sp def test_get_tokenizer_str(): - tokenizer = get_tokenizer(pretrained_model_name_or_path="test-model-name") + tokenizer = FeatureExtractor(pretrained_model_name_or_path="test-model-name") tokenizer.mocker.from_pretrained.assert_called_with( pretrained_model_name_or_path="test-model-name", revision=None, use_fast=True, use_auth_token=None ) def test_get_tokenizer_path(tmp_path): - tokenizer = get_tokenizer(pretrained_model_name_or_path=tmp_path / "test-path") + tokenizer = FeatureExtractor(pretrained_model_name_or_path=tmp_path / "test-path") tokenizer.mocker.from_pretrained.assert_called_with( pretrained_model_name_or_path=str(tmp_path / "test-path"), revision=None, use_fast=True, use_auth_token=None ) def test_get_tokenizer_keep_accents(): - tokenizer = get_tokenizer(pretrained_model_name_or_path="test-model-name-albert") + tokenizer = FeatureExtractor(pretrained_model_name_or_path="test-model-name-albert") tokenizer.mocker.from_pretrained.assert_called_with( pretrained_model_name_or_path="test-model-name-albert", revision=None, @@ -100,7 +100,7 @@ def test_get_tokenizer_keep_accents(): def test_get_tokenizer_mlm_warning(caplog): - tokenizer = get_tokenizer(pretrained_model_name_or_path="test-model-name-mlm") + tokenizer = FeatureExtractor(pretrained_model_name_or_path="test-model-name-mlm") tokenizer.mocker.from_pretrained.assert_called_with( pretrained_model_name_or_path="test-model-name-mlm", revision=None, use_fast=True, use_auth_token=None ) @@ -115,7 +115,7 @@ def test_get_tokenizer_mlm_warning(caplog): @pytest.mark.integration @pytest.mark.parametrize("model_name", TOKENIZERS_TO_TEST) def test_save_load(tmp_path, model_name: str): - tokenizer = get_tokenizer(pretrained_model_name_or_path=model_name, do_lower_case=False) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=model_name, do_lower_case=False) text = "Some Text with neverseentokens plus !215?#. and a combined-token_with/chars" tokenizer.add_tokens(new_tokens=["neverseentokens"]) @@ -124,7 +124,7 @@ def test_save_load(tmp_path, model_name: str): save_dir = tmp_path / "saved_tokenizer" tokenizer.save_pretrained(save_dir) - tokenizer_loaded = get_tokenizer(pretrained_model_name_or_path=save_dir) + tokenizer_loaded = FeatureExtractor(pretrained_model_name_or_path=save_dir) new_encoding = tokenizer_loaded(text) assert original_encoding == new_encoding @@ -132,7 +132,7 @@ def test_save_load(tmp_path, model_name: str): @pytest.mark.integration def test_tokenize_custom_vocab_bert(): - tokenizer = get_tokenizer(pretrained_model_name_or_path=BERT, do_lower_case=False) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=BERT, do_lower_case=False) tokenizer.add_tokens(new_tokens=["neverseentokens"]) text = "Some Text with neverseentokens plus !215?#. and a combined-token_with/chars" @@ -162,7 +162,7 @@ def test_tokenization_on_edge_cases_full_sequence_tokenization(model_name: str, """ Verify that tokenization on full sequence is the same as the one on "whitespace tokenized words" """ - tokenizer = get_tokenizer(pretrained_model_name_or_path=model_name, do_lower_case=False, add_prefix_space=True) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=model_name, do_lower_case=False, add_prefix_space=True) pre_tokenizer = WhitespaceSplit() words_and_spans = pre_tokenizer.pre_tokenize_str(edge_case) @@ -182,7 +182,7 @@ def test_tokenization_on_edge_cases_full_sequence_tokenization_roberta_exception Verify that tokenization on full sequence is the same as the one on "whitespace tokenized words". These test cases work for all tokenizers under test except for RoBERTa. """ - tokenizer = get_tokenizer(pretrained_model_name_or_path=model_name, do_lower_case=False, add_prefix_space=True) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=model_name, do_lower_case=False, add_prefix_space=True) pre_tokenizer = WhitespaceSplit() words_and_spans = pre_tokenizer.pre_tokenize_str(edge_case) @@ -211,7 +211,7 @@ def test_tokenization_on_edge_cases_full_sequence_tokenization_roberta_exception ) @pytest.mark.parametrize("model_name,marker", TOKENIZERS_TO_TEST_WITH_TOKEN_MARKER) def test_tokenization_on_edge_cases_full_sequence_verify_spans(model_name: str, marker: str, edge_case: str): - tokenizer = get_tokenizer(pretrained_model_name_or_path=model_name, do_lower_case=False, add_prefix_space=True) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=model_name, do_lower_case=False, add_prefix_space=True) pre_tokenizer = WhitespaceSplit() words_and_spans = pre_tokenizer.pre_tokenize_str(edge_case) @@ -246,7 +246,7 @@ def test_tokenization_on_edge_cases_full_sequence_verify_spans(model_name: str, ], ) def test_detokenization_for_bert(edge_case): - tokenizer = get_tokenizer(pretrained_model_name_or_path=BERT, do_lower_case=False) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=BERT, do_lower_case=False) encoded = tokenizer(edge_case, add_special_tokens=False).encodings[0] @@ -261,7 +261,7 @@ def test_detokenization_for_bert(edge_case): @pytest.mark.integration def test_encode_plus_for_bert(): - tokenizer = get_tokenizer(pretrained_model_name_or_path=BERT, do_lower_case=False) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=BERT, do_lower_case=False) text = "Some Text with neverseentokens plus !215?#. and a combined-token_with/chars" encoded_batch = tokenizer(text) @@ -309,7 +309,7 @@ def test_encode_plus_for_bert(): @pytest.mark.integration def test_tokenize_custom_vocab_bert(): - tokenizer = get_tokenizer(pretrained_model_name_or_path=BERT, do_lower_case=False) + tokenizer = FeatureExtractor(pretrained_model_name_or_path=BERT, do_lower_case=False) tokenizer.add_tokens(new_tokens=["neverseentokens"]) text = "Some Text with neverseentokens plus !215?#. and a combined-token_with/chars" diff --git a/test/modeling/test_dpr.py b/test/modeling/test_dpr.py index 7cc6d01193..405b979ee4 100644 --- a/test/modeling/test_dpr.py +++ b/test/modeling/test_dpr.py @@ -9,15 +9,16 @@ import pytest import torch from torch.utils.data import SequentialSampler from tqdm import tqdm -from transformers import DPRQuestionEncoder +from transformers import DPRQuestionEncoder, AutoTokenizer from haystack.modeling.data_handler.dataloader import NamedDataLoader from haystack.modeling.data_handler.processor import TextSimilarityProcessor from haystack.modeling.model.biadaptive_model import BiAdaptiveModel from haystack.modeling.model.language_model import get_language_model, DPREncoder from haystack.modeling.model.prediction_head import TextSimilarityHead -from haystack.modeling.model.tokenization import get_tokenizer + from haystack.nodes.retriever.dense import DensePassageRetriever + from haystack.modeling.utils import set_all_seeds, initialize_device_settings from ..conftest import SAMPLES_PATH @@ -31,10 +32,10 @@ def test_dpr_modules(caplog=None): devices, n_gpu = initialize_device_settings(use_cuda=True) # 1.Create question and passage tokenizers - query_tokenizer = get_tokenizer( + query_tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path="facebook/dpr-question_encoder-single-nq-base", do_lower_case=True, use_fast=True ) - passage_tokenizer = get_tokenizer( + passage_tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path="facebook/dpr-ctx_encoder-single-nq-base", do_lower_case=True, use_fast=True ) @@ -363,9 +364,9 @@ def test_dpr_processor(embed_title, passage_ids, passage_attns, use_fast, num_ha ] query_tok = "facebook/dpr-question_encoder-single-nq-base" - query_tokenizer = get_tokenizer(query_tok, use_fast=use_fast) + query_tokenizer = AutoTokenizer.from_pretrained(query_tok, use_fast=use_fast) passage_tok = "facebook/dpr-ctx_encoder-single-nq-base" - passage_tokenizer = get_tokenizer(passage_tok, use_fast=use_fast) + passage_tokenizer = AutoTokenizer.from_pretrained(passage_tok, use_fast=use_fast) processor = TextSimilarityProcessor( query_tokenizer=query_tokenizer, passage_tokenizer=passage_tokenizer, @@ -420,9 +421,9 @@ def test_dpr_processor_empty_title(use_fast, embed_title): } query_tok = "facebook/dpr-question_encoder-single-nq-base" - query_tokenizer = get_tokenizer(query_tok, use_fast=use_fast) + query_tokenizer = AutoTokenizer.from_pretrained(query_tok, use_fast=use_fast) passage_tok = "facebook/dpr-ctx_encoder-single-nq-base" - passage_tokenizer = get_tokenizer(passage_tok, use_fast=use_fast) + passage_tokenizer = AutoTokenizer.from_pretrained(passage_tok, use_fast=use_fast) processor = TextSimilarityProcessor( query_tokenizer=query_tokenizer, passage_tokenizer=passage_tokenizer, @@ -505,9 +506,9 @@ def test_dpr_problematic(): ] query_tok = "facebook/dpr-question_encoder-single-nq-base" - query_tokenizer = get_tokenizer(query_tok) + query_tokenizer = AutoTokenizer.from_pretrained(query_tok) passage_tok = "facebook/dpr-ctx_encoder-single-nq-base" - passage_tokenizer = get_tokenizer(passage_tok) + passage_tokenizer = AutoTokenizer.from_pretrained(passage_tok) processor = TextSimilarityProcessor( query_tokenizer=query_tokenizer, passage_tokenizer=passage_tokenizer, @@ -536,9 +537,9 @@ def test_dpr_query_only(): ] query_tok = "facebook/dpr-question_encoder-single-nq-base" - query_tokenizer = get_tokenizer(query_tok) + query_tokenizer = AutoTokenizer.from_pretrained(query_tok) passage_tok = "facebook/dpr-ctx_encoder-single-nq-base" - passage_tokenizer = get_tokenizer(passage_tok) + passage_tokenizer = AutoTokenizer.from_pretrained(passage_tok) processor = TextSimilarityProcessor( query_tokenizer=query_tokenizer, passage_tokenizer=passage_tokenizer, @@ -598,9 +599,9 @@ def test_dpr_context_only(): ] query_tok = "facebook/dpr-question_encoder-single-nq-base" - query_tokenizer = get_tokenizer(query_tok) + query_tokenizer = AutoTokenizer.from_pretrained(query_tok) passage_tok = "facebook/dpr-ctx_encoder-single-nq-base" - passage_tokenizer = get_tokenizer(passage_tok) + passage_tokenizer = AutoTokenizer.from_pretrained(passage_tok) processor = TextSimilarityProcessor( query_tokenizer=query_tokenizer, passage_tokenizer=passage_tokenizer, @@ -649,9 +650,9 @@ def test_dpr_processor_save_load(tmp_path): } query_tok = "facebook/dpr-question_encoder-single-nq-base" - query_tokenizer = get_tokenizer(query_tok) + query_tokenizer = AutoTokenizer.from_pretrained(query_tok) passage_tok = "facebook/dpr-ctx_encoder-single-nq-base" - passage_tokenizer = get_tokenizer(passage_tok) + passage_tokenizer = AutoTokenizer.from_pretrained(passage_tok) processor = TextSimilarityProcessor( query_tokenizer=query_tokenizer, passage_tokenizer=passage_tokenizer, @@ -727,11 +728,11 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa # load model from model hub query_embedding_model = query_and_passage_model["query"] passage_embedding_model = query_and_passage_model["passage"] - query_tokenizer = get_tokenizer( + query_tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path=query_embedding_model ) # tokenizer class is inferred automatically query_encoder = get_language_model(pretrained_model_name_or_path=query_embedding_model) - passage_tokenizer = get_tokenizer(pretrained_model_name_or_path=passage_embedding_model) + passage_tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=passage_embedding_model) passage_encoder = get_language_model(pretrained_model_name_or_path=passage_embedding_model) processor = TextSimilarityProcessor( @@ -771,11 +772,11 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa passage_tokenizer.save_pretrained(save_dir + f"/{passage_encoder_dir}") # load model from disk - loaded_query_tokenizer = get_tokenizer( + loaded_query_tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path=Path(save_dir) / query_encoder_dir, use_fast=True ) # tokenizer class is inferred automatically loaded_query_encoder = get_language_model(pretrained_model_name_or_path=Path(save_dir) / query_encoder_dir) - loaded_passage_tokenizer = get_tokenizer( + loaded_passage_tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path=Path(save_dir) / passage_encoder_dir, use_fast=True ) loaded_passage_encoder = get_language_model(pretrained_model_name_or_path=Path(save_dir) / passage_encoder_dir) @@ -893,11 +894,13 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa loaded_passage_tokenizer.save_pretrained(save_dir + f"/{passage_encoder_dir}") # load model from disk - query_tokenizer = get_tokenizer( + query_tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path=Path(save_dir) / query_encoder_dir ) # tokenizer class is inferred automatically query_encoder = get_language_model(pretrained_model_name_or_path=Path(save_dir) / query_encoder_dir) - passage_tokenizer = get_tokenizer(pretrained_model_name_or_path=Path(save_dir) / passage_encoder_dir) + passage_tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path=Path(save_dir) / passage_encoder_dir + ) passage_encoder = get_language_model(pretrained_model_name_or_path=Path(save_dir) / passage_encoder_dir) processor = TextSimilarityProcessor( diff --git a/test/modeling/test_feature_extraction.py b/test/modeling/test_feature_extraction.py new file mode 100644 index 0000000000..705411352a --- /dev/null +++ b/test/modeling/test_feature_extraction.py @@ -0,0 +1,85 @@ +import pytest +from unittest.mock import MagicMock, patch + +import haystack +from haystack.modeling.model.feature_extraction import FeatureExtractor, FEATURE_EXTRACTORS + + +class MockedAutoTokenizer: + mocker: MagicMock = MagicMock() + + @classmethod + def from_pretrained(cls, *args, **kwargs): + cls.mocker.from_pretrained(*args, **kwargs) + return cls() + + +class MockedAutoConfig: + mocker: MagicMock = MagicMock() + model_type: str = "mocked" + + @classmethod + def from_pretrained(cls, *args, **kwargs): + cls.mocker.from_pretrained(*args, **kwargs) + return cls() + + +@pytest.fixture(autouse=True) +def mock_autotokenizer(request, monkeypatch): + # Do not patch integration tests + if "integration" in request.keywords: + return + monkeypatch.setattr( + haystack.modeling.model.feature_extraction, "FEATURE_EXTRACTORS", {"mocked": MockedAutoTokenizer} + ) + monkeypatch.setattr(haystack.modeling.model.feature_extraction, "AutoConfig", MockedAutoConfig) + monkeypatch.setattr(haystack.modeling.model.feature_extraction, "AutoTokenizer", MockedAutoTokenizer) + + +# +# Unit tests +# + + +def test_init_str(): + tokenizer = FeatureExtractor(pretrained_model_name_or_path="test-model-name") + + tokenizer.feature_extractor.mocker.from_pretrained.assert_called_with( + pretrained_model_name_or_path="test-model-name", revision=None, use_fast=True, use_auth_token=None + ) + + +def test_init_path(tmp_path): + tokenizer = FeatureExtractor(pretrained_model_name_or_path=tmp_path / "test-path") + + tokenizer.feature_extractor.mocker.from_pretrained.assert_called_with( + pretrained_model_name_or_path=str(tmp_path / "test-path"), revision=None, use_fast=True, use_auth_token=None + ) + + +# +# Integration tests +# + +FEATURE_EXTRACTORS_TO_TEST = ["bert-base-cased"] + + +@pytest.mark.integration +@pytest.mark.parametrize("model_name", FEATURE_EXTRACTORS_TO_TEST) +def test_load_modify_save_load(tmp_path, model_name: str): + + # Load base tokenizer + feature_extractor = FeatureExtractor(pretrained_model_name_or_path=model_name, do_lower_case=False) + + # Add new tokens + feature_extractor.feature_extractor.add_tokens(new_tokens=["neverseentokens"]) + + # Save modified tokenizer + save_dir = tmp_path / "saved_tokenizer" + feature_extractor.feature_extractor.save_pretrained(save_dir) + + # Load modified tokenizer + new_feature_extractor = FeatureExtractor(pretrained_model_name_or_path=save_dir) + + # Assert the new tokenizer still has the added tokens + assert len(new_feature_extractor.feature_extractor) == len(feature_extractor.feature_extractor) diff --git a/test/modeling/test_processor.py b/test/modeling/test_processor.py index 79308d80f8..6bd899d724 100644 --- a/test/modeling/test_processor.py +++ b/test/modeling/test_processor.py @@ -1,10 +1,8 @@ import logging -from pathlib import Path from transformers import AutoTokenizer from haystack.modeling.data_handler.processor import SquadProcessor -from haystack.modeling.model.tokenization import get_tokenizer from ..conftest import SAMPLES_PATH @@ -24,7 +22,7 @@ def test_dataset_from_dicts_qa_inference(caplog=None): sample_types = ["answer-wrong", "answer-offset-wrong", "noanswer", "vanilla"] for model in models: - tokenizer = get_tokenizer(pretrained_model_name_or_path=model) + tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model) processor = SquadProcessor(tokenizer, max_seq_len=256, data_dir=None) for sample_type in sample_types: @@ -251,7 +249,7 @@ def test_dataset_from_dicts_qa_labelconversion(caplog=None): sample_types = ["answer-wrong", "answer-offset-wrong", "noanswer", "vanilla"] for model in models: - tokenizer = get_tokenizer(pretrained_model_name_or_path=model) + tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model) processor = SquadProcessor(tokenizer, max_seq_len=256, data_dir=None) for sample_type in sample_types: diff --git a/test/modeling/test_processor_save_load.py b/test/modeling/test_processor_save_load.py index 154b303f70..60c6a7e057 100644 --- a/test/modeling/test_processor_save_load.py +++ b/test/modeling/test_processor_save_load.py @@ -1,8 +1,9 @@ import logging from pathlib import Path +from transformers import AutoTokenizer + from haystack.modeling.data_handler.processor import SquadProcessor -from haystack.modeling.model.tokenization import get_tokenizer from haystack.modeling.utils import set_all_seeds import torch @@ -16,7 +17,7 @@ def test_processor_saving_loading(tmp_path, caplog): set_all_seeds(seed=42) lang_model = "roberta-base" - tokenizer = get_tokenizer(pretrained_model_name_or_path=lang_model, do_lower_case=False) + tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=lang_model, do_lower_case=False) processor = SquadProcessor( tokenizer=tokenizer, diff --git a/test/nodes/test_retriever.py b/test/nodes/test_retriever.py index 4e25a8a055..5e2e7aadd1 100644 --- a/test/nodes/test_retriever.py +++ b/test/nodes/test_retriever.py @@ -1,14 +1,19 @@ +from typing import List + +import os import logging import os from math import isclose +import pytest import numpy as np import pandas as pd +from pandas.testing import assert_frame_equal +from elasticsearch import Elasticsearch +from transformers import DPRContextEncoderTokenizerFast, DPRQuestionEncoderTokenizerFast + from haystack.document_stores.base import BaseDocumentStore from haystack.document_stores.memory import InMemoryDocumentStore -import pytest -from elasticsearch import Elasticsearch - from haystack.document_stores import WeaviateDocumentStore from haystack.nodes.retriever.base import BaseRetriever from haystack.pipelines import DocumentSearchPipeline @@ -18,7 +23,7 @@ from haystack.document_stores.faiss import FAISSDocumentStore from haystack.document_stores import MilvusDocumentStore from haystack.nodes.retriever.dense import DensePassageRetriever, EmbeddingRetriever, TableTextRetriever from haystack.nodes.retriever.sparse import BM25Retriever, FilterRetriever, TfidfRetriever -from transformers import DPRContextEncoderTokenizerFast, DPRQuestionEncoderTokenizerFast +from haystack.nodes.retriever.multimodal import MultiModalRetriever from ..conftest import SAMPLES_PATH @@ -623,3 +628,179 @@ def test_es_filter_only(document_store, retriever): document_store.write_documents(docs) retrieved_docs = retriever.retrieve(query="", filters={"f1": ["0"]}) assert len(retrieved_docs) == 11 + + +# +# MultiModal +# + + +@pytest.fixture +def text_docs() -> List[Document]: + return [ + Document( + content="My name is Paul and I live in New York", + meta={ + "meta_field": "test2", + "name": "filename2", + "date_field": "2019-10-01", + "numeric_field": 5.0, + "odd_field": 0, + }, + ), + Document( + content="My name is Carla and I live in Berlin", + meta={ + "meta_field": "test1", + "name": "filename1", + "date_field": "2020-03-01", + "numeric_field": 5.5, + "odd_field": 1, + }, + ), + Document( + content="My name is Christelle and I live in Paris", + meta={ + "meta_field": "test3", + "name": "filename3", + "date_field": "2018-10-01", + "numeric_field": 4.5, + "odd_field": 1, + }, + ), + Document( + content="My name is Camila and I live in Madrid", + meta={ + "meta_field": "test4", + "name": "filename4", + "date_field": "2021-02-01", + "numeric_field": 3.0, + "odd_field": 0, + }, + ), + Document( + content="My name is Matteo and I live in Rome", + meta={ + "meta_field": "test5", + "name": "filename5", + "date_field": "2019-01-01", + "numeric_field": 0.0, + "odd_field": 1, + }, + ), + ] + + +@pytest.fixture +def table_docs() -> List[Document]: + return [ + Document( + content=pd.DataFrame( + { + "Mountain": ["Mount Everest", "K2", "Kangchenjunga", "Lhotse", "Makalu"], + "Height": ["8848m", "8,611 m", "8 586m", "8 516 m", "8,485m"], + } + ), + content_type="table", + ), + Document( + content=pd.DataFrame( + { + "City": ["Paris", "Lyon", "Marseille", "Lille", "Toulouse", "Bordeaux"], + "Population": ["13,114,718", "2,280,845", "1,873,270 ", "1,510,079", "1,454,158", "1,363,711"], + } + ), + content_type="table", + ), + Document( + content=pd.DataFrame( + { + "City": ["Berlin", "Hamburg", "Munich", "Cologne"], + "Population": ["3,644,826", "1,841,179", "1,471,508", "1,085,664"], + } + ), + content_type="table", + ), + ] + + +@pytest.fixture +def image_docs() -> List[Document]: + return [ + Document(content=str(SAMPLES_PATH / "images" / imagefile), content_type="image") + for imagefile in os.listdir(SAMPLES_PATH / "images") + ] + + +@pytest.mark.integration +def test_multimodal_text_retrieval(text_docs: List[Document]): + retriever = MultiModalRetriever( + document_store=InMemoryDocumentStore(return_embedding=True), + query_embedding_model="sentence-transformers/multi-qa-mpnet-base-dot-v1", + document_embedding_models={"text": "sentence-transformers/multi-qa-mpnet-base-dot-v1"}, + ) + retriever.document_store.write_documents(text_docs) + retriever.document_store.update_embeddings(retriever=retriever) + + results = retriever.retrieve(query="Who lives in Paris?") + assert results[0].content == "My name is Christelle and I live in Paris" + + +@pytest.mark.integration +def test_multimodal_table_retrieval(table_docs: List[Document]): + retriever = MultiModalRetriever( + document_store=InMemoryDocumentStore(return_embedding=True), + query_embedding_model="deepset/all-mpnet-base-v2-table", + document_embedding_models={"table": "deepset/all-mpnet-base-v2-table"}, + ) + retriever.document_store.write_documents(table_docs) + retriever.document_store.update_embeddings(retriever=retriever) + + results = retriever.retrieve(query="How many people live in Hamburg?") + assert_frame_equal( + results[0].content, + pd.DataFrame( + { + "City": ["Berlin", "Hamburg", "Munich", "Cologne"], + "Population": ["3,644,826", "1,841,179", "1,471,508", "1,085,664"], + } + ), + ) + + +@pytest.mark.integration +def test_multimodal_image_retrieval(image_docs: List[Document]): + retriever = MultiModalRetriever( + document_store=InMemoryDocumentStore(return_embedding=True, embedding_dim=512), + query_embedding_model="sentence-transformers/clip-ViT-B-32", + document_embedding_models={"image": "sentence-transformers/clip-ViT-B-32"}, + ) + retriever.document_store.write_documents(image_docs) + retriever.document_store.update_embeddings(retriever=retriever) + + results = retriever.retrieve(query="What's a cat?") + assert str(results[0].content) == str(SAMPLES_PATH / "images" / "cat.jpg") + + +@pytest.mark.skip("Not working yet as intended") +@pytest.mark.integration +def test_multimodal_text_image_retrieval(text_docs: List[Document], image_docs: List[Document]): + retriever = MultiModalRetriever( + document_store=InMemoryDocumentStore(return_embedding=True, embedding_dim=512), + query_embedding_model="sentence-transformers/clip-ViT-B-32", + document_embedding_models={ + "text": "sentence-transformers/clip-ViT-B-32", + "image": "sentence-transformers/clip-ViT-B-32", + }, + ) + retriever.document_store.write_documents(image_docs) + retriever.document_store.write_documents(text_docs) + retriever.document_store.update_embeddings(retriever=retriever) + + results = retriever.retrieve(query="What's Paris?") + + text_results = [result for result in results if result.content_type == "text"] + image_results = [result for result in results if result.content_type == "image"] + + assert str(image_results[0].content) == str(SAMPLES_PATH / "images" / "paris.jpg") + assert text_results[0].content == "My name is Christelle and I live in Paris" diff --git a/test/samples/images/apple.jpg b/test/samples/images/apple.jpg new file mode 100644 index 0000000000..f9023fea2c Binary files /dev/null and b/test/samples/images/apple.jpg differ diff --git a/test/samples/images/car.jpg b/test/samples/images/car.jpg new file mode 100644 index 0000000000..8e11cf386e Binary files /dev/null and b/test/samples/images/car.jpg differ diff --git a/test/samples/images/cat.jpg b/test/samples/images/cat.jpg new file mode 100644 index 0000000000..25c458ac72 Binary files /dev/null and b/test/samples/images/cat.jpg differ diff --git a/test/samples/images/galaxy.jpg b/test/samples/images/galaxy.jpg new file mode 100644 index 0000000000..c13fa64073 Binary files /dev/null and b/test/samples/images/galaxy.jpg differ diff --git a/test/samples/images/paris.jpg b/test/samples/images/paris.jpg new file mode 100644 index 0000000000..8fb3d839f0 Binary files /dev/null and b/test/samples/images/paris.jpg differ