Embedded / models /train_model.py
yamunasivan's picture
create a new file
d1e319b
raw
history blame contribute delete
951 Bytes
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
import pickle
# Load dataset
file_path = "data/sms_process_data_main.xlsx"
df = pd.read_excel(file_path)
# Prepare training data
X_train, X_test, y_train, y_test = train_test_split(df['MessageText'], df['label'], test_size=0.2, random_state=42)
# Convert text into numerical vectors
vectorizer = TfidfVectorizer()
X_train_vec = vectorizer.fit_transform(X_train)
X_test_vec = vectorizer.transform(X_test)
# Train model
model = LogisticRegression(max_iter=1000)
model.fit(X_train_vec, y_train)
# Save model and vectorizer
with open("models/logistic_regression.pkl", "wb") as model_file:
pickle.dump(model, model_file)
with open("models/vectorizer.pkl", "wb") as vec_file:
pickle.dump(vectorizer, vec_file)
print("Model trained and saved successfully!")