Update app.py

app.py

@@ -39,7 +39,7 @@ st.caption("Multi-Modal Fake News & AI Image Detection System")
 def load_text_model():
     # model_name = "Maheentouqeer1/truthguard-fake-news-detector"
     # NEW - replace with this well-trained model
-    model_name = "GonzaloA/fake-news-detector"
+    model_name = "hamzab/roberta-fake-news-classification"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(
         model_name, low_cpu_mem_usage=True
@@ -94,7 +94,10 @@ def predict_news(text, tokenizer, text_model):
     probs = torch.nn.functional.softmax(outputs.logits, dim=1)
     prediction = torch.argmax(probs).item()
     confidence = torch.max(probs).item() * 100
-    return prediction, confidence
+
+    # This model outputs: 0 = FAKE, 1 = REAL
+    label = text_model.config.id2label[prediction]
+    return label, confidence
 # ---------------------------
 # PREDICT IMAGE
 # ---------------------------
@@ -128,13 +131,13 @@ with tab1:
         if user_input.strip() == "":
             st.warning("Please enter some text")
         else:
-            pred, conf = predict_news(user_input, tokenizer, text_model)
-            if pred == 0:
+            label, conf = predict_news(user_input, tokenizer, text_model)
+
+            if label == "FAKE":
                 st.error(f"⚠️ FAKE NEWS ({conf:.2f}%)")
-            elif pred == 1:
-                st.success(f"✅ REAL NEWS ({conf:.2f}%)")
             else:
-                st.warning(f"🤔 UNCERTAIN ({conf:.2f}%)")
+                st.success(f"✅ REAL NEWS ({conf:.2f}%)")
+
             st.progress(int(conf))
 
 with tab2: