(self, model_name, gpu_id, max_memory)
| 49 | bot_id = "<bot>" |
| 50 | |
| 51 | def __init__(self, model_name, gpu_id, max_memory): |
| 52 | device = torch.device('cuda', gpu_id) # TODO: allow sending to cpu |
| 53 | |
| 54 | # recommended default for devices with > 40 GB VRAM |
| 55 | # load model onto one device |
| 56 | if max_memory is None: |
| 57 | self._model = AutoModelForCausalLM.from_pretrained( |
| 58 | model_name, torch_dtype=torch.float16, device_map="auto") |
| 59 | self._model.to(device) |
| 60 | # load the model with the given max_memory config (for devices with insufficient VRAM or multi-gpu) |
| 61 | else: |
| 62 | config = AutoConfig.from_pretrained(model_name) |
| 63 | # load empty weights |
| 64 | with init_empty_weights(): |
| 65 | model_from_conf = AutoModelForCausalLM.from_config(config) |
| 66 | |
| 67 | model_from_conf.tie_weights() |
| 68 | |
| 69 | # create a device_map from max_memory |
| 70 | device_map = infer_auto_device_map( |
| 71 | model_from_conf, |
| 72 | max_memory=max_memory, |
| 73 | no_split_module_classes=["GPTNeoXLayer"], |
| 74 | dtype="float16" |
| 75 | ) |
| 76 | # load the model with the above device_map |
| 77 | self._model = AutoModelForCausalLM.from_pretrained( |
| 78 | model_name, |
| 79 | device_map=device_map, |
| 80 | offload_folder="offload", # optional offload-to-disk overflow directory (auto-created) |
| 81 | offload_state_dict=True, |
| 82 | torch_dtype=torch.float16 |
| 83 | ) |
| 84 | self._tokenizer = AutoTokenizer.from_pretrained(model_name) |
| 85 | |
| 86 | def do_inference(self, prompt, max_new_tokens, do_sample, temperature, top_k, stream_callback=None): |
| 87 | stop_criteria = StopWordsCriteria(self._tokenizer, [self.human_id], stream_callback) |
no test coverage detected