Generate text in a streaming manner for a single prompt, optionally using an image. :param prompt: The input text prompt for generation. :param image: Optional; an image tensor to be used in conjunction with the text prompt. Shape: (channel, h, w) or (1, channel, h
(
self,
prompt: str,
image: Optional[torch.FloatTensor] = None,
max_gen_len: int = 512,
temperature: float = 0.0,
top_p: float = 0.95,
additional_stop_symbols: Iterable[str] = ()
)
| 468 | |
| 469 | @ torch.inference_mode() |
| 470 | def stream_generate( |
| 471 | self, |
| 472 | prompt: str, |
| 473 | image: Optional[torch.FloatTensor] = None, |
| 474 | max_gen_len: int = 512, |
| 475 | temperature: float = 0.0, |
| 476 | top_p: float = 0.95, |
| 477 | additional_stop_symbols: Iterable[str] = () |
| 478 | ): |
| 479 | """ |
| 480 | Generate text in a streaming manner for a single prompt, optionally using an image. |
| 481 | |
| 482 | :param prompt: The input text prompt for generation. |
| 483 | :param image: Optional; an image tensor to be used in conjunction with the text prompt. |
| 484 | Shape: (channel, h, w) or (1, channel, h, w). |
| 485 | :param max_gen_len: Maximum length for the generated text. Default is 512. |
| 486 | :param temperature: Temperature for generation randomness. Default is 0.0, |
| 487 | namely deterministic generation. |
| 488 | :param top_p: Top-p sampling probability for diverse generation. Default is 0.95. |
| 489 | :param additional_stop_symbols: Iterable of additional symbols to stop generation. |
| 490 | :return: A generator yielding dictionaries with generated text and an end-of-content flag. |
| 491 | """ |
| 492 | args = self.llma.args |
| 493 | |
| 494 | prompt_tokens = self.tokenizer.encode(prompt, bos=True, eos=False) |
| 495 | # truncate from the left. leave some space for generation. |
| 496 | max_seq_len = args.max_seq_len |
| 497 | if image is not None: |
| 498 | max_seq_len -= self.llma.image_words |
| 499 | if len(image.shape) == 4: |
| 500 | assert image.shape[0] == 1 |
| 501 | else: |
| 502 | assert len(image.shape) == 3 |
| 503 | image = image.unsqueeze(0) |
| 504 | image = image.to(list(self.parameters())[0].device) |
| 505 | |
| 506 | max_prompt_size = max_seq_len - max_gen_len |
| 507 | prompt_tokens = prompt_tokens[-max_prompt_size:] |
| 508 | |
| 509 | prompt_size = len(prompt_tokens) |
| 510 | |
| 511 | total_len = min(max_seq_len, max_gen_len + prompt_size) |
| 512 | |
| 513 | tokens = torch.full([total_len], 0).cuda().long() |
| 514 | |
| 515 | tokens[:len(prompt_tokens)] = torch.tensor(prompt_tokens).long() |
| 516 | start_pos = prompt_size |
| 517 | prev_pos = 0 |
| 518 | generate_until = start_pos |
| 519 | for cur_pos in range(start_pos, total_len): |
| 520 | logits = self.llma.forward_inference( |
| 521 | tokens[None, prev_pos:cur_pos], prev_pos, image if prev_pos == 0 else None |
| 522 | ).float() |
| 523 | if temperature > 0: |
| 524 | probs = torch.softmax(logits / temperature, dim=-1) |
| 525 | next_token = self.sample_top_p(probs, top_p) |
| 526 | else: |
| 527 | next_token = torch.argmax(logits, dim=-1) |
no test coverage detected