Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2022 The REALM authors and The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""REALM Retriever model implementation."""16 17import os18from typing import Optional, Union19 20import numpy as np21from huggingface_hub import hf_hub_download22 23from transformers import AutoTokenizer24 25from ....utils import logging, strtobool26 27 28_REALM_BLOCK_RECORDS_FILENAME = "block_records.npy"29 30 31logger = logging.get_logger(__name__)32 33 34def convert_tfrecord_to_np(block_records_path: str, num_block_records: int) -> np.ndarray:35    import tensorflow.compat.v1 as tf36 37    blocks_dataset = tf.data.TFRecordDataset(block_records_path, buffer_size=512 * 1024 * 1024)38    blocks_dataset = blocks_dataset.batch(num_block_records, drop_remainder=True)39    np_record = next(blocks_dataset.take(1).as_numpy_iterator())40 41    return np_record42 43 44class ScaNNSearcher:45    """Note that ScaNNSearcher cannot currently be used within the model. In future versions, it might however be included."""46 47    def __init__(48        self,49        db,50        num_neighbors,51        dimensions_per_block=2,52        num_leaves=1000,53        num_leaves_to_search=100,54        training_sample_size=100000,55    ):56        """Build scann searcher."""57 58        from scann.scann_ops.py.scann_ops_pybind import builder as Builder59 60        builder = Builder(db=db, num_neighbors=num_neighbors, distance_measure="dot_product")61        builder = builder.tree(62            num_leaves=num_leaves, num_leaves_to_search=num_leaves_to_search, training_sample_size=training_sample_size63        )64        builder = builder.score_ah(dimensions_per_block=dimensions_per_block)65 66        self.searcher = builder.build()67 68    def search_batched(self, question_projection):69        retrieved_block_ids, _ = self.searcher.search_batched(question_projection.detach().cpu())70        return retrieved_block_ids.astype("int64")71 72 73class RealmRetriever:74    """The retriever of REALM outputting the retrieved evidence block and whether the block has answers as well as answer75    positions."76 77        Parameters:78            block_records (`np.ndarray`):79                A numpy array which contains evidence texts.80            tokenizer ([`RealmTokenizer`]):81                The tokenizer to encode retrieved texts.82    """83 84    def __init__(self, block_records, tokenizer):85        super().__init__()86        self.block_records = block_records87        self.tokenizer = tokenizer88 89    def __call__(self, retrieved_block_ids, question_input_ids, answer_ids, max_length=None, return_tensors="pt"):90        retrieved_blocks = np.take(self.block_records, indices=retrieved_block_ids, axis=0)91 92        question = self.tokenizer.decode(question_input_ids[0], skip_special_tokens=True)93 94        text = []95        text_pair = []96        for retrieved_block in retrieved_blocks:97            text.append(question)98            text_pair.append(retrieved_block.decode())99 100        concat_inputs = self.tokenizer(101            text, text_pair, padding=True, truncation=True, return_special_tokens_mask=True, max_length=max_length102        )103        concat_inputs_tensors = concat_inputs.convert_to_tensors(return_tensors)104 105        if answer_ids is not None:106            return self.block_has_answer(concat_inputs, answer_ids) + (concat_inputs_tensors,)107        else:108            return (None, None, None, concat_inputs_tensors)109 110    @classmethod111    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *init_inputs, **kwargs):112        if os.path.isdir(pretrained_model_name_or_path):113            block_records_path = os.path.join(pretrained_model_name_or_path, _REALM_BLOCK_RECORDS_FILENAME)114        else:115            block_records_path = hf_hub_download(116                repo_id=pretrained_model_name_or_path, filename=_REALM_BLOCK_RECORDS_FILENAME, **kwargs117            )118        if not strtobool(os.environ.get("TRUST_REMOTE_CODE", "False")):119            raise ValueError(120                "This part uses `pickle.load` which is insecure and will execute arbitrary code that is "121                "potentially malicious. It's recommended to never unpickle data that could have come from an "122                "untrusted source, or that could have been tampered with. If you already verified the pickle "123                "data and decided to use it, you can set the environment variable "124                "`TRUST_REMOTE_CODE` to `True` to allow it."125            )126        block_records = np.load(block_records_path, allow_pickle=True)127 128        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)129 130        return cls(block_records, tokenizer)131 132    def save_pretrained(self, save_directory):133        # save block records134        np.save(os.path.join(save_directory, _REALM_BLOCK_RECORDS_FILENAME), self.block_records)135        # save tokenizer136        self.tokenizer.save_pretrained(save_directory)137 138    def block_has_answer(self, concat_inputs, answer_ids):139        """check if retrieved_blocks has answers."""140        has_answers = []141        start_pos = []142        end_pos = []143        max_answers = 0144 145        for input_id in concat_inputs.input_ids:146            input_id_list = input_id.tolist()147            # Check answers between two [SEP] tokens148            first_sep_idx = input_id_list.index(self.tokenizer.sep_token_id)149            second_sep_idx = first_sep_idx + 1 + input_id_list[first_sep_idx + 1 :].index(self.tokenizer.sep_token_id)150 151            start_pos.append([])152            end_pos.append([])153            for answer in answer_ids:154                for idx in range(first_sep_idx + 1, second_sep_idx):155                    if answer[0] == input_id_list[idx]:156                        if input_id_list[idx : idx + len(answer)] == answer:157                            start_pos[-1].append(idx)158                            end_pos[-1].append(idx + len(answer) - 1)159 160            if len(start_pos[-1]) == 0:161                has_answers.append(False)162            else:163                has_answers.append(True)164                if len(start_pos[-1]) > max_answers:165                    max_answers = len(start_pos[-1])166 167        # Pad -1 to max_answers168        for start_pos_, end_pos_ in zip(start_pos, end_pos):169            if len(start_pos_) < max_answers:170                padded = [-1] * (max_answers - len(start_pos_))171                start_pos_ += padded172                end_pos_ += padded173        return has_answers, start_pos, end_pos174 175 176__all__ = ["RealmRetriever"]177