Update app.py

app.py

@@ -1,77 +1,44 @@
+import gradio as gr
+from transformers import pipeline, AutoConfig, AutoModel
 import torch
-import os
-from transformers import BertConfig, AutoModelForSequenceClassification, BertTokenizerFast
 
-def save_model(model, tokenizer, out_dir: str):
-    """Saves model in Hugging Face-compatible format"""
-    os.makedirs(out_dir, exist_ok=True)
-    
-    # 1. Save full model architecture
-    model.bert.save_pretrained(out_dir)
-    
-    # 2. Save custom classifier weights with standard name
-    torch.save(model.state_dict(), os.path.join(out_dir, "model_state.bin"))
-    
-    # 3. Create compatible config
-    config = BertConfig.from_pretrained("bert-base-uncased")
-    config.update({
-        "num_labels": 2,
-        "id2label": {0: "REAL", 1: "FAKE"},
-        "hidden_dropout_prob": 0.1,
-        "classifier_dropout": 0.1,
-        "model_type": "bert-for-sequence-classification"
-    })
-    config.save_pretrained(out_dir)
-    
-    # 4. Save tokenizer
-    tokenizer.save_pretrained(out_dir)
-    print(f"✅ Model saved in HF format to {out_dir}")
+model_id = "KenLumod/ML-Fake-Real-News-Detector-Final"
 
-def load_model(model_dir: str, device=torch.device("cpu")):
-    """Loads model using AutoModelForSequenceClassification"""
-    from transformers import AutoModelForSequenceClassification
-    
-    # Load with custom config
-    model = AutoModelForSequenceClassification.from_pretrained(
-        model_dir,
-        config=BertConfig.from_pretrained(model_dir)
-    )
-    
-    # Load custom weights
-    state_dict = torch.load(os.path.join(model_dir, "model_state.bin"))
-    model.load_state_dict(state_dict)
-    
-    # Load tokenizer
-    tokenizer = BertTokenizerFast.from_pretrained(model_dir)
-    
-    model.to(device).eval()
-    print(f"✅ Model loaded from {model_dir}")
-    return model, tokenizer
+# Force reload model with updated config
+config = AutoConfig.from_pretrained(model_id)
+config.id2label = {1: "Fake News", 0: "Real News"}  # Force override
+config.label2id = {v: k for k, v in config.id2label.items()}
 
-# ------------------------
-# 7) Updated Inference
-# ------------------------
-def predict(texts, model, tokenizer, device, max_length=15):
-    """Batch prediction with proper HF formatting"""
-    encodings = tokenizer(
-        texts,
-        max_length=max_length,
-        truncation=True,
-        padding="max_length",
-        return_tensors="pt"
-    ).to(device)
-    
-    with torch.no_grad():
-        outputs = model(**encodings)
-    
-    return torch.argmax(outputs.logits, dim=1).cpu().numpy()
+# Load the model using AutoModel (ensure safetensors format is handled)
+model = AutoModel.from_pretrained(
+    model_id,
+    config=config,
+    # Specify safe_tensors=True if using safetensors
+    safetensors=True
+)
 
-# Usage
-best_model, best_tokenizer = load_model("KenLumod/ML-Fake-Real-News-Detector-Final")
+# Create the pipeline for classification
+classifier = pipeline(
+    "text-classification",
+    model=model,
+    tokenizer=model_id,
+    return_all_scores=False
+)
 
-new_texts = [
-    "NASA announces alien life discovery",
-    "Congress passes new infrastructure bill"
-]
-predictions = predict(new_texts, best_model, best_tokenizer, device)
-print("Predictions:", ["FAKE" if p == 1 else "REAL" for p in predictions])
+def classify_news(text):
+    result = classifier(text)[0]
+    return result['label']
+
+demo = gr.Interface(
+    fn=classify_news,
+    inputs=gr.Textbox(lines=6, placeholder="Enter news article here..."),
+    outputs="text",
+    title="Fake News Detector",
+    description="Classifies news articles as Fake or Real",
+    examples=[
+        ["Breaking: Scientists discover chocolate prevents aging!"],
+        ["Parliament passes new climate change legislation"]
+    ]
+)
+
+demo.launch()