Update app.py

app.py

@@ -1,30 +1,91 @@
 import torch
+from torchvision import transforms
+from PIL import Image
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import gradio as gr
 
-# Load tokenizer and model from current directory
-model_path = "./DeBERTa"  # assuming all files are in the root
+# ---------------------------------------------------------
+# 1. Load DeBERTa text model
+# ---------------------------------------------------------
+model_path = "./DeBERTa"  
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForSequenceClassification.from_pretrained(model_path)
-model.eval()
+text_model = AutoModelForSequenceClassification.from_pretrained(model_path)
+text_model.eval()
 
-# Prediction function
-def predict(text):
+# ---------------------------------------------------------
+# 2. Load ViT image model
+# ---------------------------------------------------------
+class ViTModel(torch.nn.Module):
+    def __init__(self, base_model):
+        super().__init__()
+        self.model = base_model
+
+    def forward(self, x):
+        return self.model(x)
+
+# Load your ViT model weights
+vit_model = torch.load("trained_vit_final.pth", map_location=torch.device("cpu"))
+vit_model.eval()
+
+# Image preprocessing (modify if needed)
+image_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+# ---------------------------------------------------------
+# 3. Prediction functions
+# ---------------------------------------------------------
+
+def predict_text(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = text_model(**inputs)
         probs = torch.softmax(outputs.logits, dim=1).squeeze().tolist()
-    
     return {"Real News": probs[0], "Fake News": probs[1]}
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=predict,
+def predict_image(img):
+    img = Image.fromarray(img)
+    img_tensor = image_transforms(img).unsqueeze(0)
+
+    with torch.no_grad():
+        logits = vit_model(img_tensor)
+        probs = torch.softmax(logits, dim=1).squeeze().tolist()
+
+    # EDIT LABELS depending on your ViT classes
+    return {
+        "Real News": probs[0],
+        "Fake News": probs[1]
+    }
+
+# ---------------------------------------------------------
+# 4. Create Gradio tabs: Text and Image
+# ---------------------------------------------------------
+text_tab = gr.Interface(
+    fn=predict_text,
     inputs=gr.Textbox(lines=4, placeholder="Enter news article or headline..."),
     outputs=gr.Label(num_top_classes=2),
-    title="Fake News Detection (DeBERTa)",
-    description="A model trained using DeBERTa to detect fake news."
+    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)",
+)
+
+# ---------------------------------------------------------
+# 5. Combine into one app
+# ---------------------------------------------------------
+app = gr.TabbedInterface(
+    [text_tab, image_tab],
+    ["Text Detection", "Image Detection"]
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    app.launch()