(
self,
name: str,
dataset: str,
force_base_prompt: bool = False,
attn_implementation: str = "eager",
**kwargs,
)
| 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 |
nothing calls this directly
no test coverage detected