File size: 1,076 Bytes
75ec748 |
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 |
import torch
from transformers import (
AutoTokenizer,
BitsAndBytesConfig,
AutoModelForCausalLM,
)
def model_tokenizer_factory(
model_name,
huggingface_token: str,
):
# bitsandbytes config
USE_NESTED_QUANT = True # use_nested_quant
BNB_4BIT_COMPUTE_DTYPE = "bfloat16" # bnb_4bit_compute_dtype
compute_dtype = getattr(torch, BNB_4BIT_COMPUTE_DTYPE)
bnb_configs = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=compute_dtype,
bnb_4bit_use_double_quant=USE_NESTED_QUANT,
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
token=huggingface_token,
quantization_config=bnb_configs,
load_in_8bit=False, # Since we use 4bits
trust_remote_code=True,
attn_implementation="flash_attention_2",
torch_dtype=torch.float16,
)
tokenizer = AutoTokenizer.from_pretrained(model_name, token=huggingface_token)
model.eval()
return model, tokenizer
|