Update app.py

app.py

@@ -1,16 +1,20 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 from spaces import GPU
 
-# Load model & tokenizer
-MODEL_NAME = "ubiodee/plutus_llm"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,  # Use FP16 to reduce memory usage
-    device_map="auto"  # Automatically distribute across CPU/GPU
+# Load base model and tokenizer
+BASE_MODEL_NAME = "NousResearch/Meta-Llama-3-8B"
+LORA_MODEL_NAME = "ubiodee/plutus_llm"
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME, use_fast=False)
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_NAME,
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
+# Apply LoRA weights
+model = PeftModel.from_pretrained(base_model, LORA_MODEL_NAME)
 
 # Set padding token
 if tokenizer.pad_token is None: