Update app.py

app.py

@@ -1,12 +1,13 @@
 import os
-
-# Disable CUDA visibility at the start
-os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Prevents CUDA initialization
-
 import spaces
 import gradio as gr
 from huggingface_hub import InferenceClient, login
 import time
+import traceback
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import PeftModel, PeftConfig
+import bitsandbytes
+import torch
 
 
 @spaces.GPU  # Forces GPU allocation before execution
@@ -19,47 +20,34 @@ base_model_name = "meta-llama/Llama-3.1-8B"
 # Your fine-tuned LoRA adapter (uploaded to Hugging Face)
 lora_model_name = "starnernj/Early-Christian-Church-Fathers-LLaMA-3.1-Fine-Tuned"
 
+# Login because LLaMA 3.1 8B is a gated model
+login(token=os.getenv("HuggingFaceFineGrainedReadToken"))
 
-# Function to generate responses
-def chatbot_response(user_input):
-    # Re-enable CUDA inside the function for accelerate to manage
-    os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Adjust based on ZeroGPU setup
-    
-    import traceback
-    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-    from peft import PeftModel, PeftConfig
-    import bitsandbytes
-    from accelerate import Accelerator
-    import torch
-
-    # Check if CUDA is already initialized (for debugging)
-    if torch.cuda.is_initialized():
-        print("CUDA was already initialized before Accelerator!")
+# Enable 4-bit Quantization with BitsAndBytes
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,  # ✅ Enables 4-bit quantization for memory efficiency
+    bnb_4bit_compute_dtype=torch.float16,  # ✅ Uses float16 for performance
+    bnb_4bit_use_double_quant=True,  # ✅ Optimizes quantization
+    bnb_4bit_quant_type="nf4"  # ✅ Normalized Float-4 for better accuracy
+)
 
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    quantization_config=quantization_config,
+    device_map="auto"
+)
 
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 
-    accelerator = Accelerator()
+# Load LoRA Adapter
+print("Loading LoRA adapter...")
+model = PeftModel.from_pretrained(base_model, lora_model_name)
 
-    # Login because LLaMA 3.1 8B is a gated model
-    login(token=os.getenv("HuggingFaceFineGrainedReadToken"))
-    
-    model = AutoModelForCausalLM.from_pretrained(
-        base_model_name,
-        device_map="cpu",
-        torch_dtype=torch.float32  # Avoid any GPU-related dtype defaults
-    )
-    
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-    
-    
-    # Load LoRA adapter
-    # model = PeftModel.from_pretrained(model, lora_model_name)
-    
-    model = accelerator.prepare(model)
-    
+# Function to generate responses
+def chatbot_response(user_input):
     try:
-        inputs = tokenizer(user_input, return_tensors="pt").to(device)
+        inputs = tokenizer(user_input, return_tensors="pt").to("cuda")
         outputs = model.generate(**inputs, max_length=200)
         return tokenizer.decode(outputs[0], skip_special_tokens=True)
     except Exception as e:
@@ -77,4 +65,5 @@ interface = gr.Interface(
     description="A chatbot using a fine-tuned LoRA adapter on LLaMA 3.1 8B, tuned on thousands of writings of the early Christian Church Fathers.",
 )
 
-interface.launch()
+if __name__ == "__main__":
+    interface.launch()