Method generate

accessory/model/meta.py:372–467

Generate text responses based on input prompts, optionally using images and controlling generation parameters. :param prompts: A batched list of string prompts for text generation. :param images: Optional; batched image tensor to be used in conjunction with the text example

(
        self,
        prompts: List[str],
        images: Optional[torch.FloatTensor] = None,
        max_gen_len: int = 512,
        temperature: float = 0.0,
        top_p: float = 0.95,
        additional_stop_symbols: Iterable[str] = ()
    )

    @ torch.inference_mode()
    def generate(
        self,
        prompts: List[str],
        images: Optional[torch.FloatTensor] = None,
        max_gen_len: int = 512,
        temperature: float = 0.0,
        top_p: float = 0.95,
        additional_stop_symbols: Iterable[str] = ()
    ) -> List[str]:
        """
        Generate text responses based on input prompts, optionally using images and controlling generation parameters.

        :param prompts: A batched list of string prompts for text generation.
        :param images: Optional; batched image tensor to be used in conjunction with the text examples.
         Shape: (bsz, channel, h, w).
        :param max_gen_len: Maximum generation length for the responses. Default is 512.
        :param temperature: Controls randomness in generation. Higher values lead to more random outputs.
         Default is 0.0, namely deterministic generation.
        :param top_p: Top-p sampling probability for more diverse generation. Default is 0.95.
        :param additional_stop_symbols: Iterable of additional symbols to stop generation.
        :return: A list of generated text responses corresponding to each input prompt.
        """

        if isinstance(prompts, str):
            raise ValueError(f"{self.__class__}.generate expects a batched LIST of prompts, but str is given")

        if images is not None:
            images = images.to(list(self.parameters())[0].device)

        bsz = len(prompts)
        args = self.llma.args
        assert bsz <= args.max_batch_size, (bsz, args.max_batch_size)

        prompt_tokens = [self.tokenizer.encode(
            x, bos=True, eos=False) for x in prompts]

        min_prompt_size = min([len(t) for t in prompt_tokens])
        max_prompt_size = max([len(t) for t in prompt_tokens])

        max_seq_len = args.max_seq_len
        if images is not None:
            max_seq_len -= self.llma.image_words

        total_len = min(max_seq_len, max_gen_len + max_prompt_size)
        for k, t in enumerate(prompt_tokens):
            prompt_tokens[k] = t[-(max_seq_len-max_gen_len):]

        tokens = torch.full((bsz, total_len), 0).cuda().long()
        input_text_mask = torch.full((bsz, total_len), False).cuda()
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t).long()
            input_text_mask[k, : len(t)] = True
        start_pos = min_prompt_size
        prev_pos = 0

        l_stop_tokens = [[self.tokenizer.eos_id]]
        l_stop_tokens += [self.tokenizer.encode_segment(_) for _ in additional_stop_symbols]
        l_stop_tokens += [self.tokenizer.encode_wo_prefix_space(_) for _ in additional_stop_symbols]