from gradio import themes
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download

#Model Definition
class AGNewsClassifier(nn.Module):
    def __init__(self, num_classes = 4, dropout = 0.3):
        super().__init__()
        self.bert = AutoModel.from_pretrained('distilbert-base-uncased')
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(768, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids = input_ids, attention_mask = attention_mask)
        cls_output = outputs.last_hidden_state[:,0,:]
        dropped = self.dropout(cls_output)
        return self.classifier(dropped)


#declare label values (AG news) 
LABELS = {0: 'World News', 1: 'Sports News', 2:'Business News',3: 'Science & Technology News'}

#load tokenizer
tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')

#load model from the class
model = AGNewsClassifier()

load_weights = hf_hub_download(
    repo_id = 'sushamarangarajan/ag-news-classifier',
    filename = 'best_model.pt'
)

model.load_state_dict(torch.load(load_weights, map_location = 'cpu'))
model.eval()

#check if model got loaded
print('Model loaded')

#Prediction
def news_classification(text):
    # check for empty texts
    if not text.strip():
        return {}
    
    #if not empty, proceed
    tokens = tokenizer(
        text,
        max_length = 128,
        padding='max_length',
        truncation = True,
        return_tensors = 'pt'
    )


    with torch.no_grad():
        #torch.no_grad disables the gradient calculation operation temporarily
        logits = model(tokens['input_ids'],tokens['attention_mask'])
    
    probs = F.softmax(logits, dim = 1).squeeze()
    
    #return dict with labels and probability scores
    return {LABELS[i]:float(probs[i]) for i in range(4)}


# Building Gradio UI 

examples = [
["NBA player fined $25K for tossing ball into crowd 'with force"],
["NASA discovers potential signs of ancient water on Mars surface"],
["Pakistan declared open war with Afghanistan following strikes"],
["Paramount raises its bid to $31 per share for Warner Bros. Discovery"],
["Real Madrid faces Man City in the Champions League round of 16"],
["AI is greatly improving gene editing by 90%"],
]

demo = gr.Interface(
    fn = news_classification,
    inputs = gr.Textbox(
        lines = 3,
        placeholder = 'Paste a news headline or text snippet here....',
        label = 'News Text'
    ),
    outputs = gr.Label(
        label = 'Category Prediction'
    ),
    title = 'AG News Classifier',
    description = 'Fine tuned DistilBert model that classifies news text / headlines into World, Sports, Business, Science & Technology. Model was trained on 120k article data and achieved an accuracy of 94.33%',
    examples = examples,
    theme = gr.themes.Soft()
)
if __name__ == '__main__':
    demo.launch()