| 12 | |
| 13 | |
| 14 | class HuggingFaceDecoder(DecoderBase): |
| 15 | def __init__( |
| 16 | self, |
| 17 | name: str, |
| 18 | dataset: str, |
| 19 | force_base_prompt: bool = False, |
| 20 | attn_implementation: str = "eager", |
| 21 | **kwargs, |
| 22 | ): |
| 23 | super().__init__(name=name, **kwargs) |
| 24 | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 25 | |
| 26 | kwargs = { |
| 27 | "device_map": "auto", |
| 28 | "trust_remote_code": self.trust_remote_code, |
| 29 | "torch_dtype": getattr(torch, self.dtype), |
| 30 | "attn_implementation": attn_implementation, # "eager", "flash_attention_2", "sdpa" |
| 31 | } |
| 32 | self.skip_special_tokens = True |
| 33 | |
| 34 | print(f"{kwargs = }") |
| 35 | |
| 36 | self.force_base_prompt = force_base_prompt |
| 37 | self.tokenizer = AutoTokenizer.from_pretrained(name, use_fast=False) |
| 38 | if self.is_direct_completion(): # no chat template |
| 39 | self.eos += extra_eos_for_direct_completion(dataset) |
| 40 | else: # with chat template |
| 41 | self.eos += ["\n```\n"] |
| 42 | |
| 43 | print(f"{self.eos = }") |
| 44 | self.model = AutoModelForCausalLM.from_pretrained(name, **kwargs) |
| 45 | self.model = self.model.to(self.device) |
| 46 | |
| 47 | def is_direct_completion(self) -> bool: |
| 48 | return self.force_base_prompt or self.tokenizer.chat_template is None |
| 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 |