Function sample_sequence_batch

codegeex/torch/inference.py:217–326

(
        model,
        tokenizer,
        context_tokens,
        context_lengths,
        attention_mask,
        position_ids,
        seq_length,
        out_seq_length,
        maxlen=None,
        return_scores: bool = False,
        prompt_length: int = None,
        bad_ids: List = None,
        temperature: float = 1.0,
        topp: float = 1.0,
        topk: int = 0.0,
        recompute: bool = False,
        greedy: bool = False,
)

def sample_sequence_batch(
        model,
        tokenizer,
        context_tokens,
        context_lengths,
        attention_mask,
        position_ids,
        seq_length,
        out_seq_length,
        maxlen=None,
        return_scores: bool = False,
        prompt_length: int = None,
        bad_ids: List = None,
        temperature: float = 1.0,
        topp: float = 1.0,
        topk: int = 0.0,
        recompute: bool = False,
        greedy: bool = False,
):
    model.eval()
    with torch.no_grad():
        context_length = context_lengths.min().item()
        eos_id = tokenizer.eos_token_id

        counter = 0
        org_context_length = context_length

        layer_past = None
        batch_size = context_tokens.size(0)
        is_done = torch.zeros([batch_size]).byte().cuda()
        tokens = context_tokens
        if maxlen is None:
            maxlen = seq_length - 1
            if maxlen > (org_context_length + out_seq_length):
                maxlen = org_context_length + out_seq_length

        lengths = torch.ones([batch_size]).long().cuda() * maxlen
        if return_scores:
            scores = torch.zeros([batch_size]).float().cuda()

        while context_length <= (maxlen):

            if recompute:
                logits = model(tokens,
                               position_ids,
                               attention_mask,
                               prompt_length=prompt_length,
                               context_length=context_length,
                               )
                logits = logits[:, context_length - 1, :]
            else:
                if counter == 0:
                    tokens2use = tokens[:, :context_length]
                    positions2use = position_ids[:, :context_length]
                else:
                    tokens2use = tokens[:, context_length - 1].view(
                        batch_size, -1)
                    positions2use = position_ids[:, context_length - 1].view(