Update app.py

app.py

@@ -1,31 +1,13 @@
-from transformers import AutoModel, AutoTokenizer, AutoConfig
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-import torch.nn as nn
 
-# 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)
+# 1. Load your already-working model directly
+model = AutoModelForSequenceClassification.from_pretrained("KenLumod/ML-Fake-Real-News-Detector-Final")
+tokenizer = AutoTokenizer.from_pretrained("KenLumod/ML-Fake-Real-News-Detector-Final")
 
-# 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 = 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, 🙂)  # CLS token
-
-# 4. Create complete model
-model = FakeNewsClassifier(bert).eval()
+# 2. Prediction function (unchanged from your working version)
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return "FAKE" if outputs.logits.argmax() == 1 else "REAL"  # Your existing logic