File size: 1,367 Bytes
035246c 20b9178 035246c |
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 |
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch
# Configure quantization parameters
quantization_config = BitsAndBytesConfig(
load_in_4bit=True, # Load the model weights in 4-bit precision
bnb_4bit_compute_dtype=torch.bfloat16, # Use bfloat16 for computation
bnb_4bit_quant_type="nf4", # Use "nf4" quantization type
bnb_4bit_use_double_quant=True, # Enable double quantization
)
# Define the model name and path for the quantized model
model_name = "./Llama-3.1-Nemotron-Nano-8B-v1-bnb-4bit"
# Load the quantized model with the specified configuration
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=quantization_config,
device_map="auto" # Automatically allocate devices
)
# Load the tokenizer associated with the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Determine the device where the model is located
device = model.device
# Prepare input text and move it to the same device as the model
input_text = "Once upon a time"
inputs = tokenizer(input_text, return_tensors="pt").to(device)
# Perform inference
with torch.no_grad():
outputs = model.generate(**inputs, max_length=50)
# Decode the generated text
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)
|