Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import torch
-from transformers import LlamaTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 from huggingface_hub import login
 
@@ -14,17 +14,18 @@ if not hf_token:
 login(token=hf_token)
 
 # Load Model & Tokenizer
-model_name = "meta-llama/LLaMA-2-7b-chat-hf"  # Use the chat version for better instruction-following
+model_name = "meta-llama/Llama-2-7b-chat-hf"  # Use the chat model
 
 @st.cache_resource
 def load_model():
     # Load tokenizer
-    tokenizer = LlamaTokenizer.from_pretrained(model_name, token=hf_token)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
 
-    # Load model for causal language modeling
+    # Load model with 4-bit quantization on CPU
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        device_map="auto",  # Automatically maps model to available devices
+        device_map="cpu",  # Force CPU usage
+        load_in_4bit=True,  # Enable 4-bit quantization
         token=hf_token
     )
 
@@ -56,8 +57,8 @@ def classify_text(text, classes):
     return predicted_class
 
 # Streamlit UI
-st.title("📝 Text Classification with LLaMA 2")
-st.write("Powered by LLaMA 2 & Hugging Face")
+st.title("📝 Text Classification with LLaMA 2 Chat (CPU, 4-bit)")
+st.write("Powered by LLaMA 2 Chat & Hugging Face")
 
 # User Input
 user_input = st.text_area("Enter the text to classify:")
@@ -77,4 +78,4 @@ if st.button("Classify"):
         st.warning("Please enter some text to classify.")
 
 st.markdown("---")
-st.write("🔍 This app classifies text using the LLaMA 2 model.")
+st.write("🔍 This app classifies text using the LLaMA 2 Chat model with 4-bit quantization on CPU.")