import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer
model_name = "Aakash22134/fake_news_DistilBert"  # Your model repo

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
print("Model loaded successfully!")

def predict_news(title, text):
    """
    Predict if news is real or fake
    """
    # Combine title and text
    full_text = f"{title}. {text}"
    
    # Tokenize
    inputs = tokenizer(
        full_text, 
        return_tensors="pt", 
        truncation=True, 
        max_length=512,
        padding=True
    )
    
    # Predict
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class = torch.argmax(predictions, dim=1).item()
        confidence = predictions[0][predicted_class].item()
    
    # Map to labels
    label = "Real News ✓" if predicted_class == 0 else "Fake News ✗"
    
    # Create detailed output
    real_prob = predictions[0][0].item()
    fake_prob = predictions[0][1].item()
    
    return {
        "Prediction": label,
        "Confidence": f"{confidence:.2%}",
        "Real News Probability": f"{real_prob:.2%}",
        "Fake News Probability": f"{fake_prob:.2%}"
    }

# Example articles
examples = [
    [
        "Scientists Discover Water on Mars",
        "NASA scientists have confirmed the presence of liquid water on Mars, marking a major breakthrough in space exploration."
    ],
    [
        "Obama Signs Executive Order Banning The Pledge Of Allegiance",
        "In a stunning move, the President has banned the Pledge of Allegiance in all public schools effective immediately."
    ],
    [
        "New Study Shows Benefits of Exercise",
        "Researchers at Harvard University found that regular exercise can reduce the risk of heart disease by 30 percent."
    ]
]

# Create Gradio interface
demo = gr.Interface(
    fn=predict_news,
    inputs=[
        gr.Textbox(label="News Title", placeholder="Enter the news headline..."),
        gr.Textbox(
            label="News Content", 
            placeholder="Enter the article text...",
            lines=5
        )
    ],
    outputs=gr.JSON(label="Analysis Results"),
    title="🔍 Fake News Detector",
    description="""
    This AI model analyzes news articles to detect potential misinformation.
    Enter a news headline and article text to check if it's likely real or fake.
    
    **Note**: This is an AI prediction tool and should not be the sole method for determining news authenticity.
    """,
    examples=examples,
    theme=gr.themes.Soft(),
    article="""
    ### About This Model
    - Based on DistilBERT (distilbert-base-uncased)
    - Test Accuracy: 96.36%
    - Trained on 18,281 news articles
    
    ### How to Use
    1. Enter a news headline in the first box
    2. Paste the article content in the second box
    3. Click Submit to analyze
    
    ### Limitations
    - Works best on English news articles
    - May not perform well on very recent events
    - Should be used as one of multiple verification methods
    """
)

if __name__ == "__main__":
    demo.launch()