Update app.py

app.py

@@ -3,6 +3,13 @@ import nltk
 from nltk.sentiment.vader import SentimentIntensityAnalyzer
 from transformers import pipeline
 
+# 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")
+
+
+
 def greet(name):
     return "Hello " + name + "!!"
 
@@ -11,16 +18,18 @@ def classify(text):
 
 def predict_sentiment(text, model):
     if model == "finiteautomata/bertweet-base-sentiment-analysis":
-        pipe = pipeline("text-classification", model="finiteautomata/bertweet-base-sentiment-analysis")
-        out = pipe(text, return_all_scores=True)
+        out = sentiment_pipe(text, return_all_scores=True)
         return {pred["label"]: pred["score"] for pred in out[0]}
     elif model == "vader":
         nltk.download('vader_lexicon')
         sia = SentimentIntensityAnalyzer()
         return sia.polarity_scores(text)
+    elif model == "custom logistic regression":
+        x = lr_vectorizer.transform([text])
+        pred = lr_model.predict_proba(x)[0]
+        return {"neg": pred[0], "pos": pred[1]}
         
 
-
 demo = gr.Blocks()
 
 with demo: