(
self,
device="cuda",
pretrained="gpt2",
revision="main",
low_cpu_mem_usage=None,
subfolder=None,
tokenizer=None,
batch_size=1,
max_batch_size=512,
max_length=None,
load_in_8bit: Optional[bool] = False,
trust_remote_code: Optional[bool] = False,
dtype: Optional[Union[str, torch.dtype]]="auto",
)
| 21 | _DEFAULT_MAX_LENGTH = 2048 |
| 22 | |
| 23 | def __init__( |
| 24 | self, |
| 25 | device="cuda", |
| 26 | pretrained="gpt2", |
| 27 | revision="main", |
| 28 | low_cpu_mem_usage=None, |
| 29 | subfolder=None, |
| 30 | tokenizer=None, |
| 31 | batch_size=1, |
| 32 | max_batch_size=512, |
| 33 | max_length=None, |
| 34 | load_in_8bit: Optional[bool] = False, |
| 35 | trust_remote_code: Optional[bool] = False, |
| 36 | dtype: Optional[Union[str, torch.dtype]]="auto", |
| 37 | ): |
| 38 | super().__init__() |
| 39 | |
| 40 | |
| 41 | # Initialize model |
| 42 | if isinstance(pretrained, transformers.PreTrainedModel): |
| 43 | self.model = pretrained |
| 44 | self._device = self.model.device |
| 45 | |
| 46 | if tokenizer: |
| 47 | assert isinstance( |
| 48 | tokenizer, |
| 49 | transformers.PreTrainedTokenizer |
| 50 | ) or isinstance( |
| 51 | tokenizer, |
| 52 | transformers.PreTrainedTokenizerFast |
| 53 | ) |
| 54 | self.tokenizer = tokenizer |
| 55 | else: |
| 56 | # Get tokenizer |
| 57 | model_name = self.model.name_or_path |
| 58 | self.tokenizer = transformers.AutoTokenizer.from_pretrained( |
| 59 | model_name, |
| 60 | revision=revision, |
| 61 | trust_remote_code=trust_remote_code, |
| 62 | ) |
| 63 | |
| 64 | elif isinstance(pretrained, str): |
| 65 | |
| 66 | # Initialize device |
| 67 | assert isinstance(device, str) |
| 68 | device_list = set( |
| 69 | ["cuda", "cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())] |
| 70 | ) |
| 71 | if device and device in device_list: |
| 72 | self._device = torch.device(device) |
| 73 | print(f"Using device '{device}'") |
| 74 | else: |
| 75 | print("Device not specified") |
| 76 | print(f"Cuda Available? {torch.cuda.is_available()}") |
| 77 | self._device = ( |
| 78 | torch.device("cuda") |
| 79 | if torch.cuda.is_available() |
| 80 | else torch.device("cpu") |
nothing calls this directly
no test coverage detected