Method generate

trinity/common/models/tinker_model.py:51–121

Generate a responses from a prompt in async.

(self, prompt: str, **kwargs)

        )

    async def generate(self, prompt: str, **kwargs) -> Sequence[Experience]:
        """Generate a responses from a prompt in async."""
        if self.tokenizer is None:
            await self._initialize_tokenizer()

        returned_seq, is_valid = self._handle_prompt_truncation(prompt, **kwargs)
        if not is_valid:
            return returned_seq  # is_valid is False: returned_seq is a list of dummy experiences
        token_ids = returned_seq  # is_valid is True: returned_seq is prompt's token_ids

        with_chat_completion = kwargs.get("with_chat_completion", False)
        if with_chat_completion:
            create_time = int(time.time())
        output = await self._generate_internal(prompt={"prompt_token_ids": token_ids}, **kwargs)
        logprobs = kwargs.get("logprobs", self.config.logprobs)
        return_logprobs = logprobs is not None and logprobs is not False
        experiences = [
            Experience(
                tokens=torch.tensor(token_ids + sequence.tokens, dtype=torch.int32),
                logprobs=(
                    torch.tensor(sequence.logprobs, dtype=torch.float32)
                    if return_logprobs
                    else torch.tensor([], dtype=torch.float32)
                ),
                prompt_length=len(token_ids),
                prompt_text=self.tokenizer.decode(token_ids),
                response_text=self.tokenizer.decode(sequence.tokens),
            )
            for sequence in output.sequences
        ]
        if with_chat_completion:
            from openai.types.chat.chat_completion import (
                ChatCompletion,
                ChatCompletionMessage,
                ChatCompletionTokenLogprob,
                Choice,
                ChoiceLogprobs,
            )

            return_token_ids = kwargs.get("return_token_ids", False)
            chat_completion = ChatCompletion(
                id="",
                choices=[
                    Choice(
                        finish_reason=sequence.stop_reason,
                        index=i,
                        logprobs=ChoiceLogprobs(
                            content=[
                                ChatCompletionTokenLogprob(
                                    token=self.tokenizer.decode(token_id),
                                    logprob=logprob,
                                    top_logprobs=[],
                                )
                                for token_id, logprob in zip(sequence.tokens, sequence.logprobs)
                            ]
                        ),
                        message=ChatCompletionMessage(
                            content=self.tokenizer.decode(sequence.tokens), role="assistant"