add: remote_hub_from_pretrained

Dockerfile

@@ -16,7 +16,7 @@ EXPOSE 5000
 
 # Set the environment variable for Flask
 ENV FLASK_APP=api_server.py
-ENV TRANSFORMERS_CACHE=$(pwd)/.transformers_cache
+ENV TRANSFORMERS_CACHE=transformers_cache
 
 # Run the Flask application
 CMD ["flask", "run", "--host=0.0.0.0"]

api_server.py

@@ -1,7 +1,6 @@
 import os
 import time
 import numpy as np
-from huggingface_hub import hf_hub_download
 
 # Disable tensorflow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
@@ -9,15 +8,25 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 from tensorflow import keras
 from flask import Flask, jsonify, request
 
-LOCAL = False
+load_type = ''
+"""
+local,
+remote_hub_download - /cache error even using TRANSFORMERS_CACHE to root folder
+remote_hub_from_pretrained
+"""
+
+REPO_ID = "1vash/mnist_demo_model"
 
 # Load the saved model into memory
-if LOCAL is True:
+if load_type == 'local':
     model = keras.models.load_model('artifacts/models/mnist_model.h5')
-else:
-    REPO_ID = "1vash/mnist_demo_model"
-    FILENAME = "saved_model.pb"
-    model = keras.models.load_model(hf_hub_download(repo_id=REPO_ID, filename=FILENAME))
+elif load_type == 'remote_hub_download':
+    from huggingface_hub import hf_hub_download
+    model = keras.models.load_model(hf_hub_download(repo_id=REPO_ID, filename="saved_model.pb"))
+elif load_type == 'remote_hub_from_pretrained':
+    from huggingface_hub import from_pretrained_keras
+    model = from_pretrained_keras(REPO_ID)
+
 
 # Initialize the Flask application
 app = Flask(__name__)

model.py

@@ -1,46 +0,0 @@
-import os
-import random
-import numpy as np
-
-# disable tensorflow warnings
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
-
-import tensorflow as tf
-from tensorflow import keras
-from keras.datasets import mnist
-
-# Set the random seed for reproducibility, remember these lines :)
-SEED = 42
-random.seed(SEED)
-np.random.seed(SEED)
-tf.random.set_seed(SEED)
-
-# Load the dataset from keras.datasets (so noone would need to download it manually from any sources)
-(x_train, y_train), (x_test, y_test) = mnist.load_data()
-
-# Preprocess the dataset
-x_train = x_train.astype('float32') / 255.0
-x_test = x_test.astype('float32') / 255.0
-
-# Define the model architecture
-model = keras.Sequential([
-    keras.layers.Flatten(input_shape=(28, 28)),
-    keras.layers.Dense(128, activation='relu'),
-    keras.layers.Dense(10, activation='softmax')
-])
-
-# Compile and train the model
-# target in one-hot categorical_crossentropy -> [0,0,1,0,0,0,0,0,0]
-# target can be as integer sparse_categorical_crossentropy -> 3
-model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
-
-# 4-epoch is overfitting, 3-rd is okay
-model.fit(x_train, y_train, validation_data=(x_test, y_test), epochs=4, shuffle=True, batch_size=32)
-
-# Evaluate the model
-print('\n')
-_, test_accuracy = model.evaluate(x_test, y_test)
-print('Test accuracy:', test_accuracy)
-
-# Save the model
-model.save('artifacts/models/mnist_model.h5')