fixed peft in app.py

app.py

@@ -1,41 +1,25 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from peft import PeftModel
+from peft import PeftModel, PeftConfig
 
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-base_model_id = "mistralai/Mistral-7B-v0.1"
-ft_model_id = "asusevski/mistraloo-sft"
+peft_model_id = "asusevski/mistraloo-sft"
+peft_config = PeftConfig.from_pretrained(peft_model_id)
+model = AutoModelForCausalLM.from_pretrained(peft_config.base_model_name_or_path)
+model = PeftModel.from_pretrained(model, peft_model_id).to(device)
+model.eval()
 
 
 tokenizer = AutoTokenizer.from_pretrained(
-    base_model_id,
+    peft_config.base_model_name_or_path,
     add_bos_token=True
 )
 
 
-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
-)
-
-model = PeftModel.from_pretrained(base_model, ft_model_id).to(device)
-model.eval()
-
-
 def uwaterloo_output(post_title, post_text):
     prompt = f"""
 Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
@@ -67,25 +51,3 @@ iface = gr.Interface(
 
 iface.launch()
 
-
-
-
-# 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,  # Mistral, same as before
-#     quantization_config=bnb_config,  # Same quantization config as before
-#     device_map="auto",
-#     trust_remote_code=True,
-#     use_auth_token=True
-# )
-
-
-# ft_model = PeftModel.from_pretrained(base_model, "mistral-mistraloo/checkpoint-500")