transactify / LSTM_model.py
ai-venkat-r's picture
changes for uploading the model (#13)
54eddc8 verified
# LSTM_model.py
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, LSTM, Dense, Dropout
from data_preprocessing import preprocess_data, split_data
import joblib # To save the tokenizer and label encoder
# Define the LSTM model
def build_lstm_model(vocab_size, embedding_dim=64, max_len=10, lstm_units=128, dropout_rate=0.2, output_units=6):
model = Sequential()
model.add(Embedding(input_dim=vocab_size, output_dim=embedding_dim, input_length=max_len))
model.add(LSTM(units=lstm_units, return_sequences=False))
model.add(Dropout(dropout_rate))
model.add(Dense(units=output_units, activation='softmax'))
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
return model
# Main function to execute the training process
def main():
# Path to your data file
data_path = r"E:\transactify\transactify\transactify\transactify\transactify\data_set\transaction_data.csv"
# Preprocess the data
sequences, labels, tokenizer, label_encoder = preprocess_data(data_path)
# Check if preprocessing succeeded
if sequences is not None:
print("Data preprocessing successful!")
# Split the data into training and testing sets
X_train, X_test, y_train, y_test = split_data(sequences, labels)
print(f"Training data shape: {X_train.shape}, Training labels shape: {y_train.shape}")
print(f"Testing data shape: {X_test.shape}, Testing labels shape: {y_test.shape}")
# Build the LSTM model
vocab_size = tokenizer.num_words + 1 # +1 for padding token
model = build_lstm_model(vocab_size, max_len=10, output_units=len(label_encoder.classes_))
# Train the model
model.fit(X_train, y_train, epochs=50, batch_size=8, validation_data=(X_test, y_test))
# Evaluate the model
loss, accuracy = model.evaluate(X_test, y_test)
print(f"Test Loss: {loss:.4f}, Test Accuracy: {accuracy:.4f}")
# Save the model
model.save('transactify.h5')
print("Model saved as 'transactify.h5'")
# Save the tokenizer and label encoder
joblib.dump(tokenizer, 'tokenizer.joblib')
joblib.dump(label_encoder, 'label_encoder.joblib')
print("Tokenizer and LabelEncoder saved as 'tokenizer.joblib' and 'label_encoder.joblib'")
else:
print("Data preprocessing failed.")
# Execute the main function
if __name__ == "__main__":
main()