codingutf8 2022 the realm s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license realm retriever model implementation import os from typing import optional union import numpy as np from huggingfacehub import hfhubdownload from import autotokenizer from utils import logging realmblockrecordsfilename blockrecords npy logger logging getloggername def converttfrecordtonpblockrecordspath str numblockrecords int np ndarray import tensorflow compat v1 as tf blocksdataset tf data tfrecorddatasetblockrecordspath buffersize512 1024 1024 blocksdataset blocksdataset batchnumblockrecords dropremaindertrue nprecord nextblocksdataset take1 asnumpyiterator return nprecord class scannsearcher build scann searcher from scann scannops py scannopspybind import builder as builder builder builderdbdb numneighborsnumneighbors distancemeasuredotproduct builder builder tree numleavesnumleaves numleavestosearchnumleavestosearch trainingsamplesizetrainingsamplesize builder builder scoreahdimensionsperblockdimensionsperblock self searcher builder build def searchbatchedself questionprojection retrievedblockids self searcher searchbatchedquestionprojection detach cpu return retrievedblockids astypeint64 class realmretriever def initself blockrecords tokenizer super init self blockrecords blockrecords self tokenizer tokenizer def callself retrievedblockids questioninputids answerids maxlengthnone returntensorspt retrievedblocks np takeself blockrecords indicesretrievedblockids axis0 question self tokenizer decodequestioninputids0 skipspecialtokenstrue text textpair for retrievedblock in retrievedblocks text appendquestion textpair appendretrievedblock decode concatinputs self tokenizer text textpair paddingtrue truncationtrue returnspecialtokensmasktrue maxlengthmaxlength concatinputstensors concatinputs converttotensorsreturntensors if answerids is not none return self blockhasanswerconcatinputs answerids concatinputstensors else return none none none concatinputstensors classmethod def frompretrainedcls pretrainedmodelnameorpath optionalunionstr os pathlike initinputs kwargs if os path isdirpretrainedmodelnameorpath blockrecordspath os path joinpretrainedmodelnameorpath realmblockrecordsfilename else blockrecordspath hfhubdownload repoidpretrainedmodelnameorpath filenamerealmblockrecordsfilename kwargs blockrecords np loadblockrecordspath allowpickletrue tokenizer autotokenizer frompretrainedpretrainedmodelnameorpath initinputs kwargs return clsblockrecords tokenizer def savepretrainedself savedirectory save block records np saveos path joinsavedirectory realmblockrecordsfilename self blockrecords save tokenizer self tokenizer savepretrainedsavedirectory def blockhasanswerself concatinputs answerids check answers between two sep tokens pad 1 to maxanswers coding utf 8 2022 the realm s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license realm retriever model implementation note that scannsearcher cannot currently be used within the model in future versions it might however be included build scann searcher the retriever of realm outputting the retrieved evidence block and whether the block has answers as well as answer positions parameters block_records np ndarray a numpy array which cantains evidence texts tokenizer realmtokenizer the tokenizer to encode retrieved texts save block records save tokenizer check if retrieved_blocks has answers check answers between two sep tokens pad 1 to max_answers
import os
from typing import Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ... import AutoTokenizer
from ...utils import logging


_REALM_BLOCK_RECORDS_FILENAME = "block_records.npy"


logger = logging.get_logger(__name__)


def convert_tfrecord_to_np(block_records_path: str, num_block_records: int) -> np.ndarray:
    import tensorflow.compat.v1 as tf

    blocks_dataset = tf.data.TFRecordDataset(block_records_path, buffer_size=512 * 1024 * 1024)
    blocks_dataset = blocks_dataset.batch(num_block_records, drop_remainder=True)
    np_record = next(blocks_dataset.take(1).as_numpy_iterator())

    return np_record


class ScaNNSearcher:
    

    def __init__(
        self,
        db,
        num_neighbors,
        dimensions_per_block=2,
        num_leaves=1000,
        num_leaves_to_search=100,
        training_sample_size=100000,
    ):
        

        from scann.scann_ops.py.scann_ops_pybind import builder as Builder

        builder = Builder(db=db, num_neighbors=num_neighbors, distance_measure="dot_product")
        builder = builder.tree(
            num_leaves=num_leaves, num_leaves_to_search=num_leaves_to_search, training_sample_size=training_sample_size
        )
        builder = builder.score_ah(dimensions_per_block=dimensions_per_block)

        self.searcher = builder.build()

    def search_batched(self, question_projection):
        retrieved_block_ids, _ = self.searcher.search_batched(question_projection.detach().cpu())
        return retrieved_block_ids.astype("int64")


class RealmRetriever:
    

    def __init__(self, block_records, tokenizer):
        super().__init__()
        self.block_records = block_records
        self.tokenizer = tokenizer

    def __call__(self, retrieved_block_ids, question_input_ids, answer_ids, max_length=None, return_tensors="pt"):
        retrieved_blocks = np.take(self.block_records, indices=retrieved_block_ids, axis=0)

        question = self.tokenizer.decode(question_input_ids[0], skip_special_tokens=True)

        text = []
        text_pair = []
        for retrieved_block in retrieved_blocks:
            text.append(question)
            text_pair.append(retrieved_block.decode())

        concat_inputs = self.tokenizer(
            text, text_pair, padding=True, truncation=True, return_special_tokens_mask=True, max_length=max_length
        )
        concat_inputs_tensors = concat_inputs.convert_to_tensors(return_tensors)

        if answer_ids is not None:
            return self.block_has_answer(concat_inputs, answer_ids) + (concat_inputs_tensors,)
        else:
            return (None, None, None, concat_inputs_tensors)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *init_inputs, **kwargs):
        if os.path.isdir(pretrained_model_name_or_path):
            block_records_path = os.path.join(pretrained_model_name_or_path, _REALM_BLOCK_RECORDS_FILENAME)
        else:
            block_records_path = hf_hub_download(
                repo_id=pretrained_model_name_or_path, filename=_REALM_BLOCK_RECORDS_FILENAME, **kwargs
            )
        block_records = np.load(block_records_path, allow_pickle=True)

        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)

        return cls(block_records, tokenizer)

    def save_pretrained(self, save_directory):
        
        np.save(os.path.join(save_directory, _REALM_BLOCK_RECORDS_FILENAME), self.block_records)
        
        self.tokenizer.save_pretrained(save_directory)

    def block_has_answer(self, concat_inputs, answer_ids):
        
        has_answers = []
        start_pos = []
        end_pos = []
        max_answers = 0

        for input_id in concat_inputs.input_ids:
            input_id_list = input_id.tolist()
            
            first_sep_idx = input_id_list.index(self.tokenizer.sep_token_id)
            second_sep_idx = first_sep_idx + 1 + input_id_list[first_sep_idx + 1 :].index(self.tokenizer.sep_token_id)

            start_pos.append([])
            end_pos.append([])
            for answer in answer_ids:
                for idx in range(first_sep_idx + 1, second_sep_idx):
                    if answer[0] == input_id_list[idx]:
                        if input_id_list[idx : idx + len(answer)] == answer:
                            start_pos[-1].append(idx)
                            end_pos[-1].append(idx + len(answer) - 1)

            if len(start_pos[-1]) == 0:
                has_answers.append(False)
            else:
                has_answers.append(True)
                if len(start_pos[-1]) > max_answers:
                    max_answers = len(start_pos[-1])

        
        for start_pos_, end_pos_ in zip(start_pos, end_pos):
            if len(start_pos_) < max_answers:
                padded = [-1] * (max_answers - len(start_pos_))
                start_pos_ += padded
                end_pos_ += padded
        return has_answers, start_pos, end_pos