Update train.py

train.py

@@ -3,7 +3,7 @@ import tensorflow as tf
 from tensorflow.keras.preprocessing.image import ImageDataGenerator
 from tensorflow.keras.applications import VGG16
 from tensorflow.keras.layers import Flatten, Dense
-from huggingface_hub import push_to_hub_keras
+from huggingface_hub import push_to_hub
 
 # Environment variable for Hugging Face token
 sac = os.getenv('accesstoken')
@@ -65,7 +65,7 @@ model.compile(loss='binary_crossentropy',
 # Train the model
 history = model.fit(
     train_generator,
-    epochs=3,  # Adjust number of epochs based on dataset size and validation performance
+    epochs=1,  # Adjust number of epochs based on dataset size and validation performance
     validation_data=validation_generator
 )
 
@@ -81,10 +81,15 @@ test_loss, test_acc = model.evaluate(test_generator)
 print('Test accuracy:', test_acc)
 
 # Save the model for future use (optional)
-model.save('cat_dog_classifier.keras')
+# Not recommended for Hugging Face Hub upload (use tf.saved_model.save())
+
+# Export the model for Hugging Face Hub using tf.saved_model.save()
+export_dir = 'saved_model'  # Create a directory for the SavedModel
+tf.saved_model.save(model, export_dir)
+
 # Upload the model to your Hugging Face space repository
-push_to_hub_keras(
-    model,
+push_to_hub(
+    model_path=export_dir,  # Point to the SavedModel directory
     repo_id="okeowo1014/catsanddogs",
     commit_message="cats and dog image classifier with transfer learning",
     tags=["image-classifier", "data-augmentation", "class-weights"],