Update app.py

app.py

@@ -37,7 +37,9 @@ st.caption("Multi-Modal Fake News & AI Image Detection System")
 # ---------------------------
 @st.cache_resource
 def load_text_model():
-    model_name = "Maheentouqeer1/truthguard-fake-news-detector"
+    # model_name = "Maheentouqeer1/truthguard-fake-news-detector"
+    # NEW - replace with this well-trained model
+    model_name = "GonzaloA/fake-news-detector"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(
         model_name, low_cpu_mem_usage=True
@@ -64,10 +66,24 @@ def load_image_model():
     model.classifier[1] = nn.Linear(model.last_channel, 1)
     model.eval()
     return model
-
-# ---------------------------
-# PREDICT TEXT
-# ---------------------------
+    
+
+# # ---------------------------
+# # PREDICT TEXT
+# # ---------------------------
+# def predict_news(text, tokenizer, text_model):
+#     inputs = tokenizer(
+#         text, return_tensors="pt",
+#         truncation=True, padding=True, max_length=512
+#     )
+#     with torch.no_grad():
+#         outputs = text_model(**inputs)
+#     probs = torch.nn.functional.softmax(outputs.logits, dim=1)
+#      # ADD THIS to see raw scores
+#     st.write(f"Debug → Label 0 score: {probs[0][0]:.2f}, Label 1 score: {probs[0][1]:.2f}")
+#     prediction = torch.argmax(probs).item()
+#     confidence = torch.max(probs).item() * 100
+#     return prediction, confidence
 def predict_news(text, tokenizer, text_model):
     inputs = tokenizer(
         text, return_tensors="pt",
@@ -76,12 +92,9 @@ def predict_news(text, tokenizer, text_model):
     with torch.no_grad():
         outputs = text_model(**inputs)
     probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-     # ADD THIS to see raw scores
-    st.write(f"Debug → Label 0 score: {probs[0][0]:.2f}, Label 1 score: {probs[0][1]:.2f}")
     prediction = torch.argmax(probs).item()
     confidence = torch.max(probs).item() * 100
     return prediction, confidence
-
 # ---------------------------
 # PREDICT IMAGE
 # ---------------------------
@@ -118,8 +131,10 @@ with tab1:
             pred, conf = predict_news(user_input, tokenizer, text_model)
             if pred == 0:
                 st.error(f"⚠️ FAKE NEWS ({conf:.2f}%)")
-            else:
+            elif pred == 1:
                 st.success(f"✅ REAL NEWS ({conf:.2f}%)")
+            else:
+                st.warning(f"🤔 UNCERTAIN ({conf:.2f}%)")
             st.progress(int(conf))
 
 with tab2: