Update app.py

app.py

@@ -1,11 +1,11 @@
 import streamlit as st
 import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import LlamaTokenizer, AutoModelForSequenceClassification
 import os
 from huggingface_hub import login
 
 # Hugging Face Authentication
-hf_token = os.getenv("HUGGINGFACE_TOKEN", "").strip()  # Remove any newline characters
+hf_token = os.getenv("HUGGINGFACE_TOKEN", "").strip()
 
 if not hf_token:
     st.error("HUGGINGFACE_TOKEN not found. Please set your Hugging Face token.")
@@ -14,48 +14,55 @@ if not hf_token:
 login(token=hf_token)
 
 # Load Model & Tokenizer
-model_name = "meta-llama/LLaMA-2-7b"  # Replace with the correct LLaMA 2 model name
+model_name = "meta-llama/Llama-2-7b-hf"  # Ensure this is a fine-tuned classification model
 
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token, use_fast=False)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
-    return tokenizer, model
+    tokenizer = LlamaTokenizer.from_pretrained(model_name, token=hf_token)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, token=hf_token)
+    
+    # Move model to GPU if available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+
+    return tokenizer, model, device
 
-tokenizer, model = load_model()
+tokenizer, model, device = load_model()
+
+# Define class labels (Update based on your dataset)
+class_labels = ["Negative", "Neutral", "Positive"]  # Modify if your model has different classes
 
 # Function to classify text
 def classify_text(user_input):
-    # Tokenize the input
-    inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True)
+    inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True).to(device)
     
-    # Get model predictions
     with torch.no_grad():
         outputs = model(**inputs)
     
-    # Get predicted class
-    predictions = torch.argmax(outputs.logits, dim=-1)
-    predicted_class = predictions.item()  # Assuming single label output
-    
-    return predicted_class
+    logits = outputs.logits
+    probabilities = torch.nn.functional.softmax(logits, dim=-1)
+    predicted_class_idx = torch.argmax(probabilities, dim=-1).item()
+
+    return class_labels[predicted_class_idx], probabilities[0].cpu().tolist()
 
 # Streamlit UI
-st.title("Text Classification with LLaMA 2")
-st.write("Powered by LLaMA 2 and Hugging Face Transformers")
+st.title("📝 Text Classification with LLaMA 2")
+st.write("Powered by LLaMA 2 & Hugging Face")
 
 # User Input
-user_input = st.text_area("Type your text here...")
+user_input = st.text_area("Enter your text for classification:")
 
 if st.button("Classify"):
     if user_input:
-        # Get classification result
-        predicted_class = classify_text(user_input)
-        
+        predicted_class, probs = classify_text(user_input)
+
         # Display result
-        st.write(f"Predicted Class: {predicted_class}")
+        st.subheader(f"Predicted Class: {predicted_class}")
+        st.write(f"Confidence Scores: {dict(zip(class_labels, probs))}")
+
     else:
         st.warning("Please enter some text to classify.")
 
-# Add a footer or additional information if needed
 st.markdown("---")
-st.write("This application uses LLaMA 2 for text classification.")
+st.write("🔍 This app classifies text using a fine-tuned LLaMA 2 model.")
+