Fix classification pipeline

app.py

@@ -3,6 +3,7 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipe
 from peft import PeftModel
 import torch
 
+# Load model and tokenizer with adapter
 @st.cache_resource
 def load_model():
     base_model = "Qwen/Qwen3-0.6B"
@@ -19,19 +20,32 @@ def load_model():
     model = PeftModel.from_pretrained(base, adapter_path)
     model = model.merge_and_unload()
 
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    # ✅ Text classification pipeline
+    pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
     return pipe
 
-generator = load_model()
+# Load pipeline once
+classifier = load_model()
 
+# Streamlit UI
 st.title("📰 Fake News Detection")
-text = st.text_area("Enter a statement (news or claim):", height=200)
+text = st.text_area("Enter a news statement or claim:", height=200)
 
 if st.button("Check"):
     with st.spinner("Analyzing..."):
-        prompt = f"Determine if the following statement is fake or real:\n\n{text}\n\nAnswer:"
-        output = generator(prompt, max_length=50, do_sample=False)[0]['generated_text']
-        answer = output.split("Answer:")[-1].strip().split()[0]
-
+        # Get classification result
+        result = classifier(text)[0]
+        label = result['label']
+        score = result['score']
+
+        # Optional: format label nicely
+        if "1" in label or "POSITIVE" in label.upper():
+            verdict = "Real"
+            emoji = "✅"
+        else:
+            verdict = "Fake"
+            emoji = "❌"
+
+        # Show result
         st.subheader("Prediction")
-        st.success(f"🧠 The statement is likely: **{answer}**")
+        st.success(f"{emoji} The statement is likely: **{verdict}** (confidence: {score:.2f})")