update/changed model loader

app.py

@@ -1,19 +1,25 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
+from peft import AutoPeftModelForCausalLM
 import torch
 
 # Load the model and tokenizer
 def load_model():
-    base_model_name = "unsloth/llama-3.2-1b-instruct-bnb-4bit"  # Replace with your base model name
+    # base_model_name = "unsloth/llama-3.2-1b-instruct-bnb-4bit"  # Replace with your base model name
     lora_model_name = "sreyanghosh/lora_model"  # Replace with your LoRA model path
-    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        base_model_name, 
-        device_map="auto" if torch.cuda.is_available() else None,
-        load_in_8bit=not torch.cuda.is_available(),
+    # tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     base_model_name, 
+    #     device_map="auto" if torch.cuda.is_available() else None,
+    #     load_in_8bit=not torch.cuda.is_available(),
+    # )
+    # model = PeftModel.from_pretrained(model, lora_model_name)
+    
+    model = AutoPeftModelForCausalLM.from_pretrained(
+        lora_model_name, # YOUR MODEL YOU USED FOR TRAINING
+        load_in_4bit = True, # False
     )
-    model = PeftModel.from_pretrained(model, lora_model_name)
+    tokenizer = AutoTokenizer.from_pretrained(lora_model_name)
     model.eval()
     return tokenizer, model
 

requirements.txt

@@ -3,4 +3,4 @@ gradio
 transformers
 peft
 torch
-bitsandbytes==0.41.1
+bitsandbytes