Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2from peft import PeftModel
3import torch
4
5model_id = "vishalgimhan/uber-assistant"
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_compute_dtype=torch.bfloat16,
9 bnb_4bit_quant_type="nf4",
10 bnb_4bit_use_double_quant=True
11)
12
13base_model = AutoModelForCausalLM.from_pretrained(
14 "meta-llama/Llama-3.1-8B-Instruct",
15 quantization_config=bnb_config,
16 device_map="auto"
17)
18model = PeftModel.from_pretrained(base_model, model_id)
19tokenizer = AutoTokenizer.from_pretrained(model_id)