Update app.py

app.py

@@ -1,10 +1,7 @@
 import torch
-from PIL import Image
 from transformers import (
     AutoTokenizer,
-    AutoModelForSequenceClassification,
-    ViTForImageClassification,
-    AutoImageProcessor
+    AutoModelForSequenceClassification
 )
 import gradio as gr
 
@@ -20,25 +17,7 @@ text_model.eval()
 
 
 # ---------------------------------------------------------
-# 2. Load ViT image model
-# ---------------------------------------------------------
-# Your model weights are local, but your folder likely does NOT contain the processor.
-# So we load a pretrained processor instead.
-processor_model_name = "google/vit-base-patch16-224-in21k"
-
-image_processor = AutoImageProcessor.from_pretrained(processor_model_name)
-
-vit_path = "./trained_vit_final"
-
-vit_model = ViTForImageClassification.from_pretrained(
-    vit_path,
-    ignore_mismatched_sizes=True
-)
-vit_model.eval()
-
-
-# ---------------------------------------------------------
-# 3. Prediction functions
+# 2. Prediction function (text only)
 # ---------------------------------------------------------
 
 def predict_text(text):
@@ -54,45 +33,20 @@ def predict_text(text):
     }
 
 
-def predict_image(img):
-    inputs = image_processor(images=img, return_tensors="pt")
-
-    with torch.no_grad():
-        outputs = vit_model(**inputs)
-        probs = torch.softmax(outputs.logits, dim=1).squeeze().tolist()
-
-    return {
-        "Real News": float(probs[0]),
-        "Fake News": float(probs[1])
-    }
-
-
 # ---------------------------------------------------------
-# 4. Gradio UI
+# 3. Gradio UI (single tab)
 # ---------------------------------------------------------
 
-text_tab = gr.Interface(
+app = gr.Interface(
     fn=predict_text,
     inputs=gr.Textbox(lines=4, placeholder="Enter news article or headline..."),
     outputs=gr.Label(num_top_classes=2),
     title="Text Fake News Detector",
 )
 
-image_tab = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Label(num_top_classes=2),
-    title="Image Fake News Detector (ViT)",
-)
-
-app = gr.TabbedInterface(
-    [text_tab, image_tab],
-    ["Text Detection", "Image Detection"]
-)
-
 
 # ---------------------------------------------------------
-# 5. Launch
+# 4. Launch
 # ---------------------------------------------------------
 
 if __name__ == "__main__":