Update train_lstm_model.py

train_lstm_model.py

@@ -47,7 +47,7 @@ scaler = MinMaxScaler(feature_range=(0, 1))
 energy_data = scaler.fit_transform(energy_data)
 print("Data normalized successfully.")
 
-# Prepare the dataset for LSTM (using 60 previous time steps to predict the next time step)
+# Prepare the dataset for LSTM
 X, y = [], []
 for i in range(60, len(energy_data)):
     X.append(energy_data[i-60:i, 0])
@@ -55,7 +55,7 @@ for i in range(60, len(energy_data)):
 X, y = np.array(X), np.array(y)
 print(f"Prepared data for LSTM: X shape = {X.shape}, y shape = {y.shape}")
 
-# Reshape X for LSTM input (samples, time steps, features)
+# Reshape X for LSTM input
 X = X.reshape((X.shape[0], X.shape[1], 1))
 print(f"Reshaped X for LSTM input: X shape = {X.shape}")
 
@@ -71,27 +71,19 @@ model.compile(optimizer=Adam(), loss='mean_squared_error')
 model.fit(X, y, epochs=10, batch_size=32, verbose=1)
 print("Model training completed.")
 
-# Ensure the models directory exists and has write permissions
-os.makedirs('models', exist_ok=True)
-if not os.access('models', os.W_OK):
-    raise PermissionError("No write permissions for the 'models' directory.")
+# Ensure the models directory exists
+model_path = 'models/lstm_energy_model'
+os.makedirs(model_path, exist_ok=True)
+if not os.access(model_path, os.W_OK):
+    raise PermissionError("No write permissions for the 'models/lstm_energy_model' directory.")
 print("Models directory created and writable.")
 
-# Remove any existing (invalid) model file
-model_path = 'models/lstm_energy_model.h5'
-if os.path.exists(model_path):
-    os.remove(model_path)
-    print(f"Removed existing (invalid) file at {model_path}")
+# Remove any existing model directory
+if os.path.exists(model_path) and os.path.isdir(model_path):
+    shutil.rmtree(model_path)
+    print(f"Removed existing directory at {model_path}")
+os.makedirs(model_path, exist_ok=True)
 
-# Save the trained model
-try:
-    model.save(model_path)
-except Exception as e:
-    raise RuntimeError(f"Failed to save model to {model_path}: {e}")
-print(f"Model saved to {model_path}")
-
-# Verify the saved model file
-if os.path.exists(model_path) and os.path.getsize(model_path) > 1024:  # Expect at least 1 KB
-    print(f"✅ LSTM model trained and saved as {model_path} (Size: {os.path.getsize(model_path) / 1024:.2f} KB)")
-else:
-    raise RuntimeError(f"Model file {model_path} is invalid or too small (Size: {os.path.getsize(model_path)} bytes).")
+# Save the trained model in SavedModel format
+model.save(model_path, save_format='tf')
+print(f"✅ LSTM model trained and saved as SavedModel at {model_path}")