Update app.py

app.py

@@ -1,8 +1,8 @@
 import streamlit as st
 import torch
-from transformers import LlamaTokenizer, AutoModelForSequenceClassification
-from huggingface_hub import login
+from transformers import LlamaTokenizer, AutoModelForCausalLM
 import os
+from huggingface_hub import login
 
 # Hugging Face Authentication
 hf_token = os.getenv("HUGGINGFACE_TOKEN", "").strip()
@@ -14,59 +14,50 @@ if not hf_token:
 login(token=hf_token)
 
 # Load Model & Tokenizer
-model_name = "meta-llama/Llama-2-7b"
+model_name = "meta-llama/LLaMA-2-7b"  # Using the specified model
 
 @st.cache_resource
 def load_model():
     # Load tokenizer
     tokenizer = LlamaTokenizer.from_pretrained(model_name, token=hf_token)
 
-    # Load model without bitsandbytes (full precision)
-    model = AutoModelForSequenceClassification.from_pretrained(
+    # Load model for causal language modeling
+    model = AutoModelForCausalLM.from_pretrained(
         model_name,
         device_map="auto",  # Automatically maps model to available devices
         token=hf_token
     )
 
-    device = "cpu"  
-
-    return tokenizer, model, device
-
-tokenizer, model, device = load_model()
+    return tokenizer, model
 
-# Define class labels 
-class_labels = ["Negative", "Neutral", "Positive"]  
+tokenizer, model = load_model()
 
-# Function to classify text
-def classify_text(user_input):
-    inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True).to(device)
+# Function to generate text (example for usage)
+def generate_text(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = model.generate(**inputs, max_length=50)  # Adjust max_length as needed
     
-    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()
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Streamlit UI
-st.title("📝 Text Classification with LLaMA 2")
+st.title("📝 Text Generation with LLaMA 2")
 st.write("Powered by LLaMA 2 & Hugging Face")
 
 # User Input
-user_input = st.text_area("Enter your text for classification:")
+user_input = st.text_area("Enter your prompt:")
 
-if st.button("Classify"):
+if st.button("Generate"):
     if user_input:
-        predicted_class, probs = classify_text(user_input)
+        generated_text = generate_text(user_input)
 
         # Display result
-        st.subheader(f"Predicted Class: {predicted_class}")
-        st.write(f"Confidence Scores: {dict(zip(class_labels, probs))}")
-
+        st.subheader("Generated Text:")
+        st.write(generated_text)
     else:
-        st.warning("Please enter some text to classify.")
+        st.warning("Please enter a prompt.")
 
 st.markdown("---")
-st.write("🔍 This app classifies text using a fine-tuned LLaMA 2 model.")
+st.write("🔍 This app generates text using the LLaMA 2 model.")
+