Update app.py

app.py

@@ -1,61 +1,31 @@
-import gradio as gr
+from transformers import AutoModel, AutoTokenizer, AutoConfig
 import torch
-from transformers import AutoModel, AutoTokenizer
+import torch.nn as nn
 
-# 1. Model Architecture (Matches Your Training)
-class BERTFakeNewsDetector(torch.nn.Module):
-    def __init__(self):
+# 1. Load base components
+model_name = "KenLumod/ML-Fake-Real-News-Detector-Final"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+config = AutoConfig.from_pretrained(model_name)
+
+# 2. Load base BERT (without classification head)
+bert = AutoModel.from_pretrained(model_name)
+
+# 3. Add your custom classifier (must match training architecture)
+class FakeNewsClassifier(nn.Module):
+    def __init__(self, bert_model):
         super().__init__()
-        self.bert = AutoModel.from_pretrained("bert-base-uncased")
-        self.classifier = torch.nn.Sequential(
-            torch.nn.Dropout(0.1),
-            torch.nn.Linear(768, 512),
-            torch.nn.ReLU(),
-            torch.nn.Linear(512, 2),
-            torch.nn.LogSoftmax(dim=1)
+        self.bert = bert_model
+        self.classifier = nn.Sequential(
+            nn.Dropout(0.1),
+            nn.Linear(768, 512),  # Match your hidden layer size
+            nn.ReLU(),
+            nn.Linear(512, config.num_labels),  # Uses config's label count
+            nn.LogSoftmax(dim=1)
         )
     
     def forward(self, input_ids, attention_mask):
         outputs = self.bert(input_ids, attention_mask=attention_mask)
-        return self.classifier(outputs.last_hidden_state[:, 0, 🙂)
-
-# 2. Loading Function (Updated for model_state.bin)
-def load_model():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    
-    # Initialize model
-    model = BERTFakeNewsDetector()
-    
-    # Load weights - using model_state.bin instead of pytorch_model.bin
-    state_dict = torch.hub.load_state_dict_from_url(
-        "https://huggingface.co/KenLumod/ML-Fake-Real-News-Detector-Final/resolve/main/model_state.bin",
-        map_location=device
-    )
-    model.load_state_dict(state_dict)
-    model.to(device).eval()
-    
-    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-    return model, tokenizer, device
-
-# 3. Gradio App
-model, tokenizer, device = load_model()
-
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256).to(device)
-    with torch.no_grad():
-        logits = model(**inputs)
-        probs = torch.exp(logits).cpu().numpy()[0]
-    return {"REAL": float(probs[0]), "FAKE": float(probs[1])}
+        return self.classifier(outputs.last_hidden_state[:, 0, 🙂)  # CLS token
 
-gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(label="News Content", lines=3, placeholder="Paste article text..."),
-    outputs=gr.Label(label="Prediction"),
-    examples=[
-        ["Breaking: Scientists discover chocolate prevents all diseases!"],
-        ["Congress passes new infrastructure bill with bipartisan support"]
-    ],
-    title="Fake News Detector",
-    description="Classifies news content using BERT",
-    allow_flagging="never"
-).launch()
+# 4. Create complete model
+model = FakeNewsClassifier(bert).eval()