Zeroshot-RedPajama

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))