lstm.py

from keras.src.layers import Bidirectional
from sklearn.metrics import f1_score
from sklearn.preprocessing import LabelEncoder
from keras.models import Sequential
from keras.layers import Embedding, LSTM, Dense, Dropout, BatchNormalization
from keras.optimizers import Adam

class LSTMClassifier:
    def __init__(self, vocab_size, embedding_dim=128, lstm_units=256, dropout_rate=0.3, learning_rate=0.0003, num_classes=10):
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.lstm_units = lstm_units
        self.dropout_rate = dropout_rate
        self.learning_rate = learning_rate
        self.num_classes = num_classes
        self.model = self._build_model()
        self.label_encoder = LabelEncoder()

    def _build_model(self):
        model = Sequential()
        model.add(Embedding(input_dim=self.vocab_size, output_dim=self.embedding_dim, input_length=500))
        model.add(Bidirectional(LSTM(128, return_sequences=True)))
        model.add(Dropout(self.dropout_rate))  # Adding dropout
        model.add(Bidirectional(LSTM(128, return_sequences=True)))
        model.add(Dropout(self.dropout_rate))
        model.add(Bidirectional(LSTM(128)))
        model.add(BatchNormalization())
        model.add(Dense(64, activation='relu'))
        model.add(Dense(self.num_classes, activation='softmax'))  # Adjusted for multi-class classification
        model.compile(optimizer=Adam(learning_rate=self.learning_rate), loss='sparse_categorical_crossentropy',
                      metrics=['accuracy'])
        return model

    def fit(self, X_train, y_train, X_test, y_test, epochs=40, batch_size=16):
        """Trains the LSTM model."""
        # Ensure labels are integer-encoded using LabelEncoder
        y_train = self.label_encoder.fit_transform(y_train)
        y_test = self.label_encoder.transform(y_test)
        # Train the model
        self.model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, validation_data=(X_test, y_test))

    def evaluate(self, X_test, y_test):
        """Evaluates the LSTM model and returns F1 score."""
        # Predict on the test set
        y_test = self.label_encoder.transform(y_test)

        y_pred = self.model.predict(X_test)
        y_pred = y_pred.argmax(axis=1)

        # Evaluate the F1 score using scikit-learn
        f1 = f1_score(y_test, y_pred, average='macro')
        print(f"LSTM Model F1 Score: {f1:.4f}")
        return f1