class fillmaskpipelinepipeline def getmaskedindexself inputids generictensor np ndarray if self framework tf maskedindex tf whereinputids self tokenizer masktokenid numpy elif self framework pt maskedindex torch nonzeroinputids self tokenizer masktokenid astuplefalse else raise valueerrorunsupported framework return maskedindex def ensureexactlyonemasktokenself inputids generictensor np ndarray maskedindex self getmaskedindexinputids numel np prodmaskedindex shape if numel 1 raise pipelineexception fillmask self model basemodelprefix fno masktoken self tokenizer masktoken found on the input def ensureexactlyonemasktokenself modelinputs generictensor if isinstancemodelinputs list for modelinput in modelinputs self ensureexactlyonemasktokenmodelinputinputids0 else for inputids in modelinputsinputids self ensureexactlyonemasktokeninputids def preprocess self inputs returntensorsnone tokenizerkwargsnone preprocessparameters dictstr generictensor if returntensors is none returntensors self framework if tokenizerkwargs is none tokenizerkwargs modelinputs self tokenizerinputs returntensorsreturntensors tokenizerkwargs self ensureexactlyonemasktokenmodelinputs return modelinputs def forwardself modelinputs modeloutputs self modelmodelinputs modeloutputsinputids modelinputsinputids return modeloutputs def postprocessself modeloutputs topk5 targetidsnone cap topk if there are targets if targetids is not none and targetids shape0 topk topk targetids shape0 inputids modeloutputsinputids0 outputs modeloutputslogits if self framework tf maskedindex tf whereinputids self tokenizer masktokenid numpy 0 outputs outputs numpy logits outputs0 maskedindex probs stablesoftmaxlogits axis1 if targetids is not none probs tf gatherndtf squeezeprobs 0 targetids reshape1 1 probs tf expanddimsprobs 0 topk tf math topkprobs ktopk values predictions topk values numpy topk indices numpy else maskedindex torch nonzeroinputids self tokenizer masktokenid astuplefalse squeeze1 fill mask pipeline supports only one masktoken per sample logits outputs0 maskedindex probs logits softmaxdim1 if targetids is not none probs probs targetids values predictions probs topktopk result singlemask values shape0 1 for i values predictions in enumeratezipvalues tolist predictions tolist row for v p in zipvalues predictions copy is important since we re going to modify this array in place tokens inputids numpy copy if targetids is not none p targetidsp tolist tokensmaskedindexi p filter padding out tokens tokensnp wheretokens self tokenizer padtokenid originally we skip special tokens to give readable output for multi masks though the other mask would be removed otherwise making the output look odd so we add them back sequence self tokenizer decodetokens skipspecialtokenssinglemask proposition score v token p tokenstr self tokenizer decodep sequence sequence row appendproposition result appendrow if singlemask return result0 return result def gettargetidsself targets topknone if isinstancetargets str targets targets try vocab self tokenizer getvocab except exception vocab targetids for target in targets id vocab gettarget none if id is none inputids self tokenizer target addspecialtokensfalse returnattentionmaskfalse returntokentypeidsfalse maxlength1 truncationtrue inputids if leninputids 0 logger warning fthe specified target token target does not exist in the model vocabulary we cannot replace it with anything meaningful ignoring it continue id inputids0 xxx if users encounter this pass it becomes pretty slow so let s make sure the warning enables them to fix the input to get faster performance logger warning fthe specified target token target does not exist in the model vocabulary freplacing with self tokenizer convertidstotokensid targetids appendid targetids listsettargetids if lentargetids 0 raise valueerrorat least one target must be provided when passed targetids np arraytargetids return targetids def sanitizeparametersself topknone targetsnone tokenizerkwargsnone preprocessparams if tokenizerkwargs is not none preprocessparamstokenizerkwargs tokenizerkwargs postprocessparams if targets is not none targetids self gettargetidstargets topk postprocessparamstargetids targetids if topk is not none postprocessparamstopk topk if self tokenizer masktokenid is none raise pipelineexception fillmask self model basemodelprefix the tokenizer does not define a masktoken return preprocessparams postprocessparams def callself inputs args kwargs outputs super callinputs kwargs if isinstanceinputs list and leninputs 1 return outputs0 return outputs top_k int defaults to 5 the number of predictions to return targets str or list str optional when passed the model will limit the scores to the passed targets instead of looking up in the whole vocab if the provided targets are not in the model vocab they will be tokenized and the first resulting token will be used with a warning and that might be slower masked language modeling prediction pipeline using any modelwithlmhead see the masked language modeling examples task_summary masked language modeling for more information example python from transformers import pipeline fill_masker pipeline model bert base uncased fill_masker this is a simple mask score 0 042 token 3291 token_str problem sequence this is a simple problem score 0 031 token 3160 token_str question sequence this is a simple question score 0 03 token 8522 token_str equation sequence this is a simple equation score 0 027 token 2028 token_str one sequence this is a simple one score 0 024 token 3627 token_str rule sequence this is a simple rule learn more about the basics of using a pipeline in the pipeline tutorial pipeline_tutorial this mask filling pipeline can currently be loaded from pipeline using the following task identifier fill mask the models that this pipeline can use are models that have been trained with a masked language modeling objective which includes the bi directional models in the library see the up to date list of available models on huggingface co models https huggingface co models filter fill mask tip this pipeline only works for inputs with exactly one token masked experimental we added support for multiple masks the returned values are raw model output and correspond to disjoint probabilities where one might expect joint probabilities see discussion https github com huggingface transformers pull 10222 tip tip this pipeline now supports tokenizer_kwargs for example try python from transformers import pipeline fill_masker pipeline model bert base uncased tokenizer_kwargs truncation true fill_masker this is a simple mask with a large amount of repeated text appended 100 tokenizer_kwargs tokenizer_kwargs tip cap top_k if there are targets fill mask pipeline supports only one mask_token per sample copy is important since we re going to modify this array in place filter padding out originally we skip special tokens to give readable output for multi masks though the other mask would be removed otherwise making the output look odd so we add them back xxx if users encounter this pass it becomes pretty slow so let s make sure the warning enables them to fix the input to get faster performance fill the masked token in the text s given as inputs args args str or list str one or several texts or one list of prompts with masked tokens targets str or list str optional when passed the model will limit the scores to the passed targets instead of looking up in the whole vocab if the provided targets are not in the model vocab they will be tokenized and the first resulting token will be used with a warning and that might be slower top_k int optional when passed overrides the number of predictions to return return a list or a list of list of dict each result comes as list of dictionaries with the following keys sequence str the corresponding input with the mask token prediction score float the corresponding probability token int the predicted token id to replace the masked one token_str str the predicted token to replace the masked one
from typing import Dict

import numpy as np

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, GenericTensor, Pipeline, PipelineException


if is_tf_available():
    import tensorflow as tf

    from ..tf_utils import stable_softmax


if is_torch_available():
    import torch


logger = logging.get_logger(__name__)


@add_end_docstrings(
    PIPELINE_INIT_ARGS,
    r,
)
class FillMaskPipeline(Pipeline):
    

    def get_masked_index(self, input_ids: GenericTensor) -> np.ndarray:
        if self.framework == "tf":
            masked_index = tf.where(input_ids == self.tokenizer.mask_token_id).numpy()
        elif self.framework == "pt":
            masked_index = torch.nonzero(input_ids == self.tokenizer.mask_token_id, as_tuple=False)
        else:
            raise ValueError("Unsupported framework")
        return masked_index

    def _ensure_exactly_one_mask_token(self, input_ids: GenericTensor) -> np.ndarray:
        masked_index = self.get_masked_index(input_ids)
        numel = np.prod(masked_index.shape)
        if numel < 1:
            raise PipelineException(
                "fill-mask",
                self.model.base_model_prefix,
                f"No mask_token ({self.tokenizer.mask_token}) found on the input",
            )

    def ensure_exactly_one_mask_token(self, model_inputs: GenericTensor):
        if isinstance(model_inputs, list):
            for model_input in model_inputs:
                self._ensure_exactly_one_mask_token(model_input["input_ids"][0])
        else:
            for input_ids in model_inputs["input_ids"]:
                self._ensure_exactly_one_mask_token(input_ids)

    def preprocess(
        self, inputs, return_tensors=None, tokenizer_kwargs=None, **preprocess_parameters
    ) -> Dict[str, GenericTensor]:
        if return_tensors is None:
            return_tensors = self.framework
        if tokenizer_kwargs is None:
            tokenizer_kwargs = {}

        model_inputs = self.tokenizer(inputs, return_tensors=return_tensors, **tokenizer_kwargs)
        self.ensure_exactly_one_mask_token(model_inputs)
        return model_inputs

    def _forward(self, model_inputs):
        model_outputs = self.model(**model_inputs)
        model_outputs["input_ids"] = model_inputs["input_ids"]
        return model_outputs

    def postprocess(self, model_outputs, top_k=5, target_ids=None):
        
        if target_ids is not None and target_ids.shape[0] < top_k:
            top_k = target_ids.shape[0]
        input_ids = model_outputs["input_ids"][0]
        outputs = model_outputs["logits"]

        if self.framework == "tf":
            masked_index = tf.where(input_ids == self.tokenizer.mask_token_id).numpy()[:, 0]

            outputs = outputs.numpy()

            logits = outputs[0, masked_index, :]
            probs = stable_softmax(logits, axis=-1)
            if target_ids is not None:
                probs = tf.gather_nd(tf.squeeze(probs, 0), target_ids.reshape(-1, 1))
                probs = tf.expand_dims(probs, 0)

            topk = tf.math.top_k(probs, k=top_k)
            values, predictions = topk.values.numpy(), topk.indices.numpy()
        else:
            masked_index = torch.nonzero(input_ids == self.tokenizer.mask_token_id, as_tuple=False).squeeze(-1)
            

            logits = outputs[0, masked_index, :]
            probs = logits.softmax(dim=-1)
            if target_ids is not None:
                probs = probs[..., target_ids]

            values, predictions = probs.topk(top_k)

        result = []
        single_mask = values.shape[0] == 1
        for i, (_values, _predictions) in enumerate(zip(values.tolist(), predictions.tolist())):
            row = []
            for v, p in zip(_values, _predictions):
                
                tokens = input_ids.numpy().copy()
                if target_ids is not None:
                    p = target_ids[p].tolist()

                tokens[masked_index[i]] = p
                
                tokens = tokens[np.where(tokens != self.tokenizer.pad_token_id)]
                
                
                
                sequence = self.tokenizer.decode(tokens, skip_special_tokens=single_mask)
                proposition = {"score": v, "token": p, "token_str": self.tokenizer.decode([p]), "sequence": sequence}
                row.append(proposition)
            result.append(row)
        if single_mask:
            return result[0]
        return result

    def get_target_ids(self, targets, top_k=None):
        if isinstance(targets, str):
            targets = [targets]
        try:
            vocab = self.tokenizer.get_vocab()
        except Exception:
            vocab = {}
        target_ids = []
        for target in targets:
            id_ = vocab.get(target, None)
            if id_ is None:
                input_ids = self.tokenizer(
                    target,
                    add_special_tokens=False,
                    return_attention_mask=False,
                    return_token_type_ids=False,
                    max_length=1,
                    truncation=True,
                )["input_ids"]
                if len(input_ids) == 0:
                    logger.warning(
                        f"The specified target token `{target}` does not exist in the model vocabulary. "
                        "We cannot replace it with anything meaningful, ignoring it"
                    )
                    continue
                id_ = input_ids[0]
                
                
                
                
                logger.warning(
                    f"The specified target token `{target}` does not exist in the model vocabulary. "
                    f"Replacing with `{self.tokenizer.convert_ids_to_tokens(id_)}`."
                )
            target_ids.append(id_)
        target_ids = list(set(target_ids))
        if len(target_ids) == 0:
            raise ValueError("At least one target must be provided when passed.")
        target_ids = np.array(target_ids)
        return target_ids

    def _sanitize_parameters(self, top_k=None, targets=None, tokenizer_kwargs=None):
        preprocess_params = {}

        if tokenizer_kwargs is not None:
            preprocess_params["tokenizer_kwargs"] = tokenizer_kwargs

        postprocess_params = {}

        if targets is not None:
            target_ids = self.get_target_ids(targets, top_k)
            postprocess_params["target_ids"] = target_ids

        if top_k is not None:
            postprocess_params["top_k"] = top_k

        if self.tokenizer.mask_token_id is None:
            raise PipelineException(
                "fill-mask", self.model.base_model_prefix, "The tokenizer does not define a `mask_token`."
            )
        return preprocess_params, {}, postprocess_params

    def __call__(self, inputs, *args, **kwargs):
        
        outputs = super().__call__(inputs, **kwargs)
        if isinstance(inputs, list) and len(inputs) == 1:
            return outputs[0]
        return outputs