Update app.py

app.py

@@ -1,8 +1,8 @@
 import streamlit as st
 import torch
 from transformers import LlamaTokenizer, AutoModelForSequenceClassification
-import os
 from huggingface_hub import login
+import os
 
 # Hugging Face Authentication
 hf_token = os.getenv("HUGGINGFACE_TOKEN", "").strip()
@@ -21,10 +21,11 @@ def load_model():
     # Load tokenizer
     tokenizer = LlamaTokenizer.from_pretrained(model_name, token=hf_token)
 
-    # Load model with 8-bit quantization
+    # Load model with bitsandbytes for 8-bit quantization
     model = AutoModelForSequenceClassification.from_pretrained(
         model_name,
-        quantization_config="8bit",  # Use "4bit" for 4-bit quantization
+        load_in_8bit=True,  # Set this for 8-bit quantization
+        device_map="auto",  # Automatically maps model to available devices
         token=hf_token
     )
 
@@ -72,4 +73,3 @@ if st.button("Classify"):
 
 st.markdown("---")
 st.write("🔍 This app classifies text using a fine-tuned LLaMA 2 model.")
-