Update app.py

app.py

@@ -2,27 +2,39 @@ import os
 import gradio as gr
 import torch
 import json
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 
 # Set Hugging Face Token for Authentication (ensure it's set in your environment)
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
+# Base model (needed for QLoRA adapter)
+BASE_MODEL = "meta-llama/Llama-3-1B-Instruct"
+QLORA_ADAPTER = "meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8"
+
 # Function to load Llama model
-def load_llama_model(model_name):
-    from transformers import LlamaForCausalLM, LlamaTokenizer
-    
-    # Use AutoTokenizer which will handle various tokenizer types
-    tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_TOKEN, use_fast=False)
-    
-    # Use the LlamaForCausalLM class which can properly load the consolidated.00.pth format
-    model = LlamaForCausalLM.from_pretrained(
-        model_name,
-        token=HUGGINGFACE_TOKEN,
-        torch_dtype=torch.float16,  # Use float16 to reduce memory usage on CPU
-        low_cpu_mem_usage=True,     # Optimize for low memory usage
-        device_map="cpu"
+def load_llama_model():
+    print("Loading base model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch.bfloat16 if torch.has_bfloat16 else torch.float32,  # Use bfloat16 if available, else float32
+        device_map="cpu",  # Ensure it runs on CPU
+        token=HUGGINGFACE_TOKEN
+    )
+
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False, token=HUGGINGFACE_TOKEN)
+
+    print("Loading QLoRA adapter...")
+    model = PeftModel.from_pretrained(
+        model,
+        QLORA_ADAPTER,
+        token=HUGGINGFACE_TOKEN
     )
     
+    print("Merging LoRA weights...")
+    model = model.merge_and_unload()  # Merge LoRA weights for inference
+
     return tokenizer, model
 
 # Load Llama 3.2 model