import gradio as gr
import torch
from transformers import BertTokenizerFast, BertForSequenceClassification

def predict_news_category(text, model, tokenizer):
    # Tokenize input text
    inputs = tokenizer(text, truncation=True, padding=True, max_length=512, return_tensors='pt')
    
    # Predict
    outputs = model(**inputs)
    probs = outputs[0].softmax(1)
    
    # Get the predicted category
    _, predicted_category = torch.max(probs, dim=1)
    
    return predicted_category.item()

# Load your model
model = BertForSequenceClassification.from_pretrained('akhil2808/EPICS-PROJECT')
model.eval()  # Set the model to evaluation mode

# Load tokenizer
tokenizer = BertTokenizerFast.from_pretrained('akhil2808/EPICS-PROJECT')

# Function to predict news category
def detect_news_category(text):
    category = predict_news_category(text, model, tokenizer)
    # Map the prediction to fake or real news
    prediction_dict = {1: 'Real News', 0: 'Fake News'}
    return prediction_dict[category]

iface = gr.Interface(fn=detect_news_category, 
                     inputs=gr.inputs.Textbox(lines=7, placeholder='News Here...'), 
                     outputs='text',
                     title='Disinformation Detector',
                     description='In the age of information, disinformation spreads rapidly. Fake news can cause substantial harm and mislead people. Therefore, it\'s crucial to detect and debunk fake news. This tool helps to detect disinformation by classifying the news as "Real" or "Fake". Powered by Group 40.',
                     theme='huggingface')

iface.launch()