File size: 1,769 Bytes
769af1a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import data
import config
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
def __train_model(df, full=False):
### train on full data for final model else split the data and then train
if full:
X_train = df["Clean_Text"]
y_train = df["Emotion"]
else:
X_train, X_test, y_train, y_test = train_test_split(
df["Clean_Text"], df["Emotion"], test_size=0.2, random_state=42
)
### build model pipeline
lr_pipeline = Pipeline(
steps=[("cv", CountVectorizer()), ("lr", LogisticRegression(max_iter=300))]
)
### train and test the model
print(f"\nTraining LogisticRegression with {X_train.shape[0]} samples...")
lr_pipeline.fit(X_train, y_train)
if not full:
print(f"Testing LogisticRegression with {X_test.shape[0]} samples...")
score = lr_pipeline.score(X_test, y_test)
print(f"Accuracy achieved: [{score*100:.2f}%].")
return lr_pipeline
if __name__ == "__main__":
emotions_df = data.load_emotions_data(config.app_config.emotions_data_file)
emotions_df = data.preprocess_data(emotions_df)
model = __train_model(emotions_df, full=True)
data.save_model(model, config.app_config.model_file)
print(f"Saved model to: [{config.app_config.model_file}]")
### Test code
# model = data.load_model(config.app_config.model_file)
# test_text = "I am loving NLP and it makes me feel so good"
# print(f"\nTesting model with sample text '{test_text}'\nPrediction:")
# print(model.predict([test_text]))
# print(model.classes_)
# print(model.predict_proba([test_text]))
|