ThinhLLM/IMDB_Sentiment_Classification_With_Bert

1 2import gradio as gr3import os4import torch5 6from model import create_bert7from timeit import default_timer as timer8from typing import Tuple, Dict9 10class_names = ['Negative','Positive']11 12model, tokenizer = create_bert()13 14model.load_state_dict(15    torch.load(16      f='bert.pth',17      map_location='cpu')18)19 20def predict(text):21    start_time = timer()22    tokenized_text = tokenizer(text,return_tensors='pt')23    24    model.eval()25    with torch.inference_mode():26        outputs = model(**tokenized_text)27    logits = outputs.logits28    probabilities = torch.softmax(logits,dim=-1)29    pred_labels_and_probs = {class_names[i]: float(probabilities[0][i]) for i in range(len(class_names))}30    pred_time = round(timer() - start_time,5)31    32    return pred_labels_and_probs, pred_time33    34    35 36title = 'IMDB Sentiment Classification 🎬🎥🍿'37des = 'A model based on BERT'38article = "Finetuned with IMDB dataset"39 40examples = [41    "I absolutely loved this movie! The acting and story were fantastic.",42    "The film was okay, but the pacing was too slow for my taste.",43    "Terrible! I wish I could get my two hours back."44]45 46demo = gr.Interface(fn=predict,47                    inputs=gr.Textbox(label="Enter your review"),48                    outputs=[gr.Label(num_top_classes=2,label='Predictions'),49                             gr.Number(label='Prediction Time (s)')],50                    examples=examples,51                    title=title,52                    description=des,53                    article=article)54demo.launch()55