| import config as config |
| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig |
| from peft import PeftModel |
| class Load_model(): |
| def __init__(self,adapter_path): |
| self.repo = config.repo_id |
| self.base_model_name = config.MODEL |
| self.adapter_path = adapter_path |
| self.token = config.HUGGINGFACE_TOKEN |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| bnb_config = BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| bnb_4bit_use_double_quant=False, |
| ) |
| base_model = AutoModelForCausalLM.from_pretrained( |
| self.base_model_name, |
| quantization_config=bnb_config, |
| device_map={"": 0}, |
| token=self.token, |
| trust_remote_code=True |
| ) |
|
|
| self.tokenizer = AutoTokenizer.from_pretrained( |
| self.base_model_name, |
| token=self.token, |
| trust_remote_code=True) |
| |
| if self.tokenizer.pad_token is None: |
| self.tokenizer.pad_token = self.tokenizer.eos_token |
| self.tokenizer.padding_side = "right" |
|
|
| self.model = PeftModel.from_pretrained(base_model, self.repo, subfolder=self.adapter_path) |
| self.model = self.model.eval() |
|
|
| print(f"--- SFT model load complete: {self.base_model_name} + {self.adapter_path} ---") |
| |
| def get(self): |
| return self.model, self.tokenizer |
|
|
