Update app.py

app.py

@@ -2,12 +2,16 @@ import gradio as gr
 import nltk
 from nltk.sentiment.vader import SentimentIntensityAnalyzer
 from transformers import pipeline
+from joblib import load
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch.nn.functional as F
 
 # global variables to load models
 lr_model = load("lr_model.joblib")
 lr_vectorizer = load("vectorizer.joblib")
 sentiment_pipe = pipeline("text-classification", model="finiteautomata/bertweet-base-sentiment-analysis")
-
+bert_model = AutoModelForSequenceClassification.from_pretrained("./imdb-bert")
+bert_tokenizer = AutoTokenizer.from_pretrained("./imdb-bert")
 
 
 def greet(name):
@@ -28,6 +32,9 @@ def predict_sentiment(text, model):
         x = lr_vectorizer.transform([text])
         pred = lr_model.predict_proba(x)[0]
         return {"neg": pred[0], "pos": pred[1]}
+    elif model == "custom BERT":
+        pred = F.softmax(bert_model(**bert_tokenizer("I love you", return_tensors="pt")).logits[0], dim=0).tolist()
+        return {"neg": pred[0], "pos": pred[1]}
         
 
 demo = gr.Blocks()