Update app.py

app.py

@@ -14,7 +14,7 @@ if not hf_token:
 login(token=hf_token)
 
 # Load Model & Tokenizer
-model_name = "meta-llama/LLaMA-2-7b"  # Use the correct model name
+model_name = "meta-llama/LLaMA-2-7b-chat-hf"  # Use the chat version for better instruction-following
 
 @st.cache_resource
 def load_model():
@@ -32,31 +32,49 @@ def load_model():
 
 tokenizer, model = load_model()
 
-# Function to generate text (example for usage)
-def generate_text(prompt):
+# Function to classify text using a prompt-based approach
+def classify_text(text, classes):
+    # Create a prompt for classification
+    prompt = f"""
+    Classify the following text into one of these categories: {", ".join(classes)}.
+    Text: {text}
+    Category:
+    """
+    
+    # Tokenize the prompt
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
+    # Generate the output
     with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=50)  # Adjust max_length as needed
+        outputs = model.generate(**inputs, max_length=100, num_return_sequences=1)
+    
+    # Decode the output
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Extract the predicted class
+    predicted_class = decoded_output.split("Category:")[-1].strip()
+    return predicted_class
 
 # Streamlit UI
-st.title("📝 Text Generation with LLaMA 2")
+st.title("📝 Text Classification with LLaMA 2")
 st.write("Powered by LLaMA 2 & Hugging Face")
 
 # User Input
-user_input = st.text_area("Enter your prompt:")
+user_input = st.text_area("Enter the text to classify:")
+
+# Define classes for classification
+classes = ["Positive", "Negative", "Neutral"]
 
-if st.button("Generate"):
+if st.button("Classify"):
     if user_input:
-        generated_text = generate_text(user_input)
+        # Perform classification
+        predicted_class = classify_text(user_input, classes)
 
         # Display result
-        st.subheader("Generated Text:")
-        st.write(generated_text)
+        st.subheader("Predicted Class:")
+        st.write(predicted_class)
     else:
-        st.warning("Please enter a prompt.")
+        st.warning("Please enter some text to classify.")
 
 st.markdown("---")
-st.write("🔍 This app generates text using the LLaMA 2 model.")
+st.write("🔍 This app classifies text using the LLaMA 2 model.")