Add TinyBERT demo

app.py

@@ -9,6 +9,7 @@ from huggingface_hub import hf_hub_download
 import torch
 import pickle
 import numpy as np
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 # Load models and tokenizers
 models = {
@@ -20,6 +21,10 @@ models = {
     "BERT Multilingual (NLP Town)": {
         "tokenizer": AutoTokenizer.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment"),
         "model": AutoModelForSequenceClassification.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment"),
+    },
+    "TinyBERT": {
+        "tokenizer": AutoTokenizer.from_pretrained("elo4/TinyBERT-sentiment-model"),
+        "model": AutoModelForSequenceClassification.from_pretrained("elo4/TinyBERT-sentiment-model"),
     }
 }
 
@@ -68,12 +73,24 @@ def predict_with_bert_multilingual(text):
         predictions = logits.argmax(axis=-1).cpu().numpy()
     return int(predictions[0] + 1)
 
+def predict_with_tinybert(text):
+    tokenizer = models["TinyBERT"]["tokenizer"]
+    model = models["TinyBERT"]["model"]
+    encodings = tokenizer([text], padding=True, truncation=True, max_length=128, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**encodings)
+        logits = outputs.logits
+        predictions = logits.argmax(axis=-1).cpu().numpy()
+    return int(predictions[0])
+
+
 # Unified function for sentiment analysis and statistics
 def analyze_sentiment_and_statistics(text):
     results = {
         "DistilBERT": predict_with_distilbert(text),
         "Logistic Regression": predict_with_logistic_regression(text),
         "BERT Multilingual (NLP Town)": predict_with_bert_multilingual(text),
+        "TinyBERT": predict_with_tinybert(text),
     }
     
     # Calculate statistics
@@ -133,7 +150,8 @@ with gr.Blocks(css=".gradio-container { max-width: 900px; margin: auto; padding:
         with gr.Column():
             distilbert_output = gr.Textbox(label="Predicted Sentiment (DistilBERT)", interactive=False)
             log_reg_output = gr.Textbox(label="Predicted Sentiment (Logistic Regression)", interactive=False)
-            bert_output = gr.Textbox(label="Predicted Sentiment (BERT Multilingual)", interactive=False)
+            bert_output = gr.Textbox(label="Predicted Sentiment (BERT Multilingual)", interactive=False),
+            tinybert_output = gr.Textbox(label="Predicted Sentiment (TinyBERT)", interactive=False)
         
         with gr.Column():
             statistics_output = gr.Textbox(label="Statistics (Lowest, Highest, Average)", interactive=False)
@@ -145,14 +163,17 @@ with gr.Blocks(css=".gradio-container { max-width: 900px; margin: auto; padding:
             f"{results['DistilBERT']}",
             f"{results['Logistic Regression']}",
             f"{results['BERT Multilingual (NLP Town)']}",
+            f"{results['TinyBERT']}",
             f"Statistics:\n{statistics['Lowest Score']}\n{statistics['Highest Score']}\nAverage Score: {statistics['Average Score']}"
         )
     
     analyze_button.click(
         process_input_and_analyze,
         inputs=[text_input],
-        outputs=[distilbert_output, log_reg_output, bert_output, statistics_output]
+        outputs=[distilbert_output, log_reg_output, bert_output, tinybert_output, statistics_output]
     )
 
+
+
 # Launch the app
 demo.launch()