@@ -52,14 +52,18 @@ def handle_classify(self, df, classifier):
52
52
#return classifier_switcher.get(classifier, lambda: "Invalid Classifier")(df)
53
53
54
54
def get_pipeline (self , model_name ):
55
+ if os .path .exists ('assets/tweet_emotions.csv' ):
56
+ prefix = ''
57
+ else :
58
+ prefix = 'public/'
55
59
if model_name == "emotion_pipeline.pkl" :
56
- df = pd .read_csv ('assets/tweet_emotions.csv' )
60
+ df = pd .read_csv (prefix + 'assets/tweet_emotions.csv' )
57
61
train_data , test_data , train_target , test_target = train_test_split (df ['content' ], df ['sentiment' ], test_size = 0.2 , shuffle = True )
58
62
elif model_name == "hate_speech.pkl" :
59
- df = pd .read_csv ('assets/nb_hatespeech.csv' , sep = ';' )
63
+ df = pd .read_csv (prefix + 'assets/nb_hatespeech.csv' , sep = ';' )
60
64
train_data , test_data , train_target , test_target = train_test_split (df ['comment' ], df ['isHate' ], test_size = 0.2 , shuffle = True )
61
65
elif model_name == "text_classification_pipeline.pkl" :
62
- df = pd .read_csv ('assets/nb_news.csv' )
66
+ df = pd .read_csv (prefix + 'assets/nb_news.csv' )
63
67
train_data , test_data , train_target , test_target = train_test_split (df ['short_description' ], df ['category' ], test_size = 0.2 , shuffle = True )
64
68
else :
65
69
with open (f'api/models/{ model_name } ' , 'rb' ) as file :
@@ -109,9 +113,12 @@ def pretrained_predict(self, df, pipeline, model_name = None):
109
113
texts_to_predict = [str (text ) for text in texts_to_predict ]
110
114
111
115
predictions = pipeline .predict (texts_to_predict )
112
- label_predictions = label_encoder .inverse_transform (predictions )
113
116
114
- df ['output_column' ] = label_predictions
117
+ if model_name :
118
+ label_predictions = label_encoder .inverse_transform (predictions )
119
+ df ['output_column' ] = label_predictions
120
+ else :
121
+ df ['output_column' ] = predictions
115
122
116
123
return df
117
124
0 commit comments