Update app.py

app.py

@@ -1,34 +1,42 @@
 import numpy as np
 import os
 import gradio as gr
-import pickle
-
 import xgboost as xgb
+import pickle
 from sklearn.feature_extraction.text import TfidfVectorizer
 
+
 os.environ["WANDB_DISABLED"] = "true"
 
+label2id =  {
+    0: "negative",
+    1: "neutral",
+    2: "positive"
+  }
+# names of the files saved in step 2: Training
+
 model_file_name = "xgb_reg.pkl"
 vectorizer_file_name = 'vectorizer.pk'
 
+
 # load
 xgb_model_loaded = pickle.load(open(model_file_name, "rb"))
 vectorizer_loaded = pickle.load(open(vectorizer_file_name, "rb"))
 
-# predict
+
 def predict_sentiment(predict_texts):
-    predictions_loaded = xgb_model_loaded.predict(vectorizer_loaded.transform(predict_texts))
-    return predictions_loaded
-    #le.inverse_transform(predictions_loaded)
+    predictions_loaded = xgb_model_loaded.predict(vectorizer_loaded.transform([predict_texts]))
+    print(predictions_loaded)
+    return label2id[predictions_loaded[0]]
 
 
 interface = gr.Interface(
     fn=predict_sentiment,
     inputs='text',
     outputs=['text'],
-    title='Croatian Movie reviews Sentiment Analysis',
+    title='Croatian Book reviews Sentiment Analysis',
     examples= ["Volim kavu","Ne volim kavu"],
     description='Get the positive/neutral/negative sentiment for the given input.'
 )
 
-interface.launch(inline = False)   
+interface.launch(inline = False)