File size: 2,068 Bytes
cfde818 4a542e2 cfde818 cb500ea cfde818 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
### 1. Imports and class names setup ###
import gradio as gr
import os
import torch
from transformers import DistilBertTokenizerFast
from timeit import default_timer as timer
# Setup class names
class_names = ["Positive", "Negative",]
### 2. Load the model ###
model = torch.load(f="BERT_sentiment_analysis.pth",
map_location=torch.device("cpu")) # load to CPU)
### 3. Predict function ###
# Create predict function
def predict(text: str):
"""Transforms and performs a prediction on img and returns prediction and time taken.
"""
# Start the timer
start_time = timer()
tokenizer = DistilBertTokenizerFast.from_pretrained(
'distilbert-base-uncased'
)
input = tokenizer(text, return_tensors="pt").to("cpu")
model.eval()
with torch.inference_mode():
logits = model(**input).logits
predicted_class_id = logits.argmax().item()
if predicted_class_id == 1:
result = "Positive π"
else:
result = "Negative π"
# Calculate the prediction time
pred_time = round(timer() - start_time, 5)
# Return the prediction dictionary and prediction time
return result
### 4. Gradio app ###
# Create title, description and article strings
title = "Sentiment Classifier"
description = "A Sentiment Classifier trained by fine-tuning [DistilBert](https://huggingface.co/docs/transformers/v4.42.0/en/model_doc/distilbert#transformers.DistilBertForSequenceClassification) Transformer model using hugging face [transformers](https://huggingface.co/docs/transformers/en/index) library."
article = "The model classifies sentiment of an input text (whether the text shows a positive or negative sentiment)."
#Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
inputs=[gr.Textbox(label="Input")],
outputs=[gr.Label(label="Prediction")],
title=title,
description=description,
article=article)
# Launch the demo!
demo.launch()
|