add config

app.py

@@ -6,9 +6,12 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 checkpoint = "Pipatpong/vcm_santa"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else "CPU"
+
+quantization_config = BitsAndBytesConfig(load_in_8bit_fp32_cpu_offload=True)
+
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True, low_cpu_mem_usage=True, load_in_8bit=True, device_map="auto", quantization_config=quantization_config)
 
 def generate(text, max_length, num_return_sequences=1):
     inputs = tokenizer.encode(text, padding=False, add_special_tokens=False, return_tensors="pt")