Update app.py

app.py

@@ -22,14 +22,19 @@ bnb_config = BitsAndBytesConfig(
     #bnb_4bit_use_double_quant=True
 )
 
-model = AutoModelForCausalLM.from_pretrained(base_model_name,
-                                             token=os.environ['HF_TOKEN'],
-                                             quantization_config=bnb_config,
-                                             device_map="auto")
-model = PeftModel.from_pretrained(model, adapter_model_name, token=os.environ['HF_TOKEN'])
-model.merge_and_unload()
-
-model = model.to(device)
+@spaces.GPU
+def load_model():
+    model = AutoModelForCausalLM.from_pretrained(base_model_name,
+                                                 token=os.environ['HF_TOKEN'],
+                                                 quantization_config=bnb_config,
+                                                 device_map="auto")
+    model = PeftModel.from_pretrained(model, adapter_model_name, token=os.environ['HF_TOKEN'])
+    model.merge_and_unload()
+    
+    model = model.to(device)
+    return model
+
+model = load_model()
 
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)