Update app.py

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+iimport streamlit as st
 import torch
 from transformers import LlamaTokenizer, AutoModelForSequenceClassification
 import os
@@ -14,35 +14,38 @@ if not hf_token:
 login(token=hf_token)
 
 # Load Model & Tokenizer
-model_name = "meta-llama/Llama-2-7b"  # Ensure this is a fine-tuned classification model
+model_name = "meta-llama/Llama-2-7b"
 
 @st.cache_resource
 def load_model():
+    # Load tokenizer
     tokenizer = LlamaTokenizer.from_pretrained(model_name, token=hf_token)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, token=hf_token)
+
+    # Load model with 8-bit quantization
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name,
+        quantization_config="8bit",  # Use "4bit" for 4-bit quantization
+        token=hf_token
+    )
 
     # Move model to GPU if available
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
 
-    # Convert model to 8-bit
-    model.half()  # Change to 16-bit for FP16, if necessary
-    model.eval()  # Set the model to evaluation mode
-
     return tokenizer, model, device
 
 tokenizer, model, device = load_model()
 
 # Define class labels (Update based on your dataset)
-class_labels = ["Negative", "Neutral", "Positive"]
+class_labels = ["Negative", "Neutral", "Positive"]  # Modify if your model has different classes
 
 # Function to classify text
 def classify_text(user_input):
     inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True).to(device)
-
+    
     with torch.no_grad():
         outputs = model(**inputs)
-
+    
     logits = outputs.logits
     probabilities = torch.nn.functional.softmax(logits, dim=-1)
     predicted_class_idx = torch.argmax(probabilities, dim=-1).item()
@@ -69,3 +72,4 @@ if st.button("Classify"):
 
 st.markdown("---")
 st.write("🔍 This app classifies text using a fine-tuned LLaMA 2 model.")
+