app.py peftmodel update

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
+from peft import PeftModel
 
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -12,10 +13,27 @@ ft_model_id = "asusevski/mistraloo-sft"
 
 tokenizer = AutoTokenizer.from_pretrained(
     base_model_id,
-    add_bos_token=True,
+    add_bos_token=True
 )
 
-model = AutoModelForCausalLM.from_pretrained(ft_model_id).to(device)
+
+base_model_id = "mistralai/Mistral-7B-v0.1"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True,
+    token=True
+)
+
+model = PeftModel.from_pretrained(base_model, ft_model_id).to(device)
 model.eval()