(
self, prompt: str, do_sample: bool = True, num_samples: int = 200
)
| 49 | |
| 50 | @torch.inference_mode() |
| 51 | def codegen( |
| 52 | self, prompt: str, do_sample: bool = True, num_samples: int = 200 |
| 53 | ) -> List[str]: |
| 54 | if self.temperature == 0: |
| 55 | assert not do_sample |
| 56 | assert num_samples == 1 |
| 57 | |
| 58 | prompt = ( |
| 59 | prompt |
| 60 | if self.is_direct_completion() |
| 61 | else make_raw_chat_prompt( |
| 62 | prompt, self.instruction_prefix, self.response_prefix, self.tokenizer |
| 63 | ) |
| 64 | ) |
| 65 | input_tokens = self.tokenizer.encode(prompt, return_tensors="pt").to( |
| 66 | self.device |
| 67 | ) |
| 68 | kwargs = {} |
| 69 | if do_sample: |
| 70 | kwargs["top_p"] = 0.95 |
| 71 | kwargs["temperature"] = self.temperature |
| 72 | |
| 73 | stop_sequencer = StopSequencer( |
| 74 | self.model, |
| 75 | model_type="causal", # or seq2seq |
| 76 | tokenizer=self.tokenizer, |
| 77 | ) |
| 78 | |
| 79 | model = stop_sequencer.register_stop_texts( |
| 80 | stop_texts=self.eos, |
| 81 | input_length=input_tokens.size(-1), |
| 82 | ) |
| 83 | |
| 84 | outputs = model.generate( |
| 85 | input_tokens, |
| 86 | max_new_tokens=self.max_new_tokens, |
| 87 | do_sample=do_sample, |
| 88 | num_return_sequences=min(self.batch_size, num_samples), |
| 89 | pad_token_id=self.tokenizer.eos_token_id, |
| 90 | **kwargs, |
| 91 | ) |
| 92 | |
| 93 | gen_strs = self.tokenizer.batch_decode( |
| 94 | outputs[:, input_tokens.size(-1) :], |
| 95 | skip_special_tokens=self.skip_special_tokens, |
| 96 | ) |
| 97 | outputs = [] |
| 98 | # removes eos tokens. |
| 99 | for output in gen_strs: |
| 100 | min_index = 10000 |
| 101 | for eos in self.eos: |
| 102 | if eos in output: |
| 103 | min_index = min(min_index, output.index(eos)) |
| 104 | outputs.append(output[:min_index].replace("\t", " ")) |
| 105 | return outputs |
nothing calls this directly
no test coverage detected