File size: 1,538 Bytes
8e83e54 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | 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
|