How to load the model using quantization:
import os
import bitsandbytes as bnb
import pandas as pd
import torch
import torch.nn as nn
import transformers
from datasets import load_dataset
from peft import (
LoraConfig,
PeftConfig,
PeftModel,
get_peft_model,
prepare_model_for_kbit_training,
)
from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
)
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
DEVICE = "cuda:0"
model_name = "Weni/ZeroShot-RedPajama"
bnb = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
config = PeftConfig.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
config.base_model_name_or_path,
quantization_config=bnb,
return_dict=True,
device_map="auto",
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token
model = PeftModel.from_pretrained(model, model_name)
_config = model.generation_config
_config.max_new_tokens = 20
_config.temperature = 0.1
_config.top_p = 0.5
_config.num_return_sequences = 1
_config.pad_token_id = tokenizer.eos_token_id
_config.eos_token_id = tokenizer.eos_token_id
Inference:
text = f"""YOUR INPUT PROMPT HERE""".strip()
encoding = tokenizer(text, return_tensors="pt").to(DEVICE)
with torch.inference_mode():
outputs = model.generate(
input_ids=encoding.input_ids,
attention_mask=encoding.attention_mask,
generation_config=_config,
do_sample=False,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))