Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import BertTokenizer, BertForSequenceClassification
+from transformers import BertTokenizer, BertForSequenceClassification,DistilBertTokenizer,DistilBertForSequenceClassification
 import torch
 import streamlit as st
 
@@ -7,6 +7,10 @@ tokenizer = BertTokenizer.from_pretrained(
 model = BertForSequenceClassification.from_pretrained(
     "ashish-001/Bert-Amazon-review-sentiment-classifier")
 
+distil_model = DistilBertForSequenceClassification.from_pretrained(
+    "ashish-001/DistilBert-Amazon-review-sentiment-classifier")
+distil_tokenizer = DistilBertTokenizer.from_pretrained(
+    "ashish-001/DistilBert-Amazon-review-sentiment-classifier")
 
 def classify_text(text):
     inputs = tokenizer(
@@ -21,10 +25,23 @@ def classify_text(text):
     probs = torch.nn.functional.sigmoid(logits)
     return probs
 
+def classify_text_distilbert(text):
+    inputs=distil_tokenizer(text, return_tensors="pt")
+    output = distil_model(**inputs)
+    logits = output.logits
+    probs = torch.nn.functional.sigmoid(logits)
+    return probs
+
 
 st.title("Amazon Review Sentiment classifier")
 data = st.text_area("Enter or paste a review")
-if st.button('Predict'):
+if st.button('Predict using BERT'):
     prediction = classify_text(data)
     st.header(
         f"Negative Confidence: {prediction[0][0].item()}, Positive Confidence: {prediction[0][1].item()}")
+
+if st.button('Predict Using DistilBERT'):
+    prediction = classify_text_distilbert(data)
+    st.header(
+        f"Negative Confidence: {prediction[0][0].item()}, Positive Confidence: {prediction[0][1].item()}")
+