Mod to support both gguf and hf model download

Dockerfile

@@ -10,10 +10,18 @@ RUN apt-get update && apt-get install -y \
 RUN useradd -m -u 1000 user
 
 # Set environment variables for the non-root user
+# Name -> Name of the app container
+# EC_ADMIN_PWD -> A secret
+# HF_MODEL_NAME -> Name of the Hugging Face Hub model
+# GGUF_MODEL_URL -> For special loading for GGUF
+# MODEL_CLASS -> A switch to load 'gguf' or 'hf'
 ENV HOME=/home/user \
     PATH=/home/user/.local/bin:$PATH \
     NAME=EduConnect \
-    EC_ADMIN_PWD='$2b$12$wGncNhE7OVmsb7TKFuNPKuJfKOIKdGtw302VMDJbAPrHrY73jqID.'
+    EC_ADMIN_PWD='$2b$12$wGncNhE7OVmsb7TKFuNPKuJfKOIKdGtw302VMDJbAPrHrY73jqID.' \
+    HF_MODEL_NAME="BitBasher/llama-2-7b-mini-ibased-GGUF" \
+    GGUF_MODEL_URL='https://huggingface.co/BitBasher/llama-2-7b-mini-ibased-GGUF/raw/main/llama-2-7b-mini-ibased.Q5_K_M.gguf' \
+    MODEL_CLASS='gguf'
 
 # Set the non-root user's home directory as the working directory
 WORKDIR $HOME
@@ -44,5 +52,14 @@ EXPOSE 7860
 # This directory is intended for persistent storage
 VOLUME /home/user/data
 
+# Copy the entrypoint script into the container and ensure it is executable
+COPY --chown=user:user entrypoint.sh $HOME
+
+# Change permission of entrypoint.sh and make sure it is executable
+RUN chmod +x $HOME/entrypoint.sh
+
+# Set the entrypoint script to be executed when the container starts
+ENTRYPOINT ["./entrypoint.sh"]
+
 # Run the FastAPI application using Uvicorn, binding to port 7860
 CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/utils/download_model.py

@@ -1,24 +1,48 @@
-# /home/user/app/utils/download_model.py
 import os
+import requests
 from transformers import AutoModel
 
-def download_model():
-    # Use environment variables for the model name and directory
+def download_hf_model():
     model_name = os.getenv("HF_MODEL_NAME")
     model_dir = f"/home/user/data/models/{model_name}"
-
+    
     # Authenticate with Hugging Face using the token, if available
     hf_token = os.getenv("HF_TOKEN")
     if hf_token:
         from huggingface_hub import HfFolder
-        HfFolder.save_token(hf_token)  # Save the token for later use by the library
-
-    # Download the model
+        HfFolder.save_token(hf_token)
+    
     print(f"Downloading model: {model_name}...")
     model = AutoModel.from_pretrained(model_name)
     model.save_pretrained(model_dir)
     print(f"Model {model_name} downloaded and saved to {model_dir}")
 
+def download_gguf_model():
+    model_name = os.getenv("HF_MODEL_NAME")
+    model_dir = f"/home/user/data/models/{model_name}"
+    os.makedirs(model_dir, exist_ok=True)
+    
+    model_url = os.getenv("GGUF_MODEL_URL")  # Assuming URL is provided as an env variable
+
+    model_file_path = os.path.join(model_dir, os.path.basename(model_url))
+    
+    print(f"Downloading model from {model_url}...")
+    response = requests.get(model_url, stream=True)
+    if response.status_code == 200:
+        with open(model_file_path, 'wb') as f:
+            f.write(response.content)
+        print(f"Model downloaded and saved to {model_file_path}")
+    else:
+        print(f"Failed to download the model. Status code: {response.status_code}")
+
+def download_model():
+    model_class = os.getenv("MODEL_CLASS")
+    if model_class == 'gguf':
+        download_gguf_model()
+    elif model_class == 'hf':
+        download_hf_model()
+    else:
+        print(f"Unsupported model class: {model_class}")
+
 if __name__ == "__main__":
     download_model()
-

entrypoint.sh

@@ -3,8 +3,8 @@
 # Authenticate with Hugging Face
 export HF_HOME=/home/user/data/hf_cache
 
+# Assuming HF_TOKEN is already exported to the environment
 echo "Using Hugging Face API token for authentication"
-export HF_TOKEN=${HF_TOKEN}
 
 # Use the environment variable for the model name
 MODEL_DIR="/home/user/data/models/${HF_MODEL_NAME}"
@@ -13,11 +13,15 @@ MODEL_URL=${HF_MODEL_URL}  # Ensure consistent variable naming
 # Download the model if it does not exist
 if [ ! -d "$MODEL_DIR" ]; then
     echo "Model not found. Downloading ${HF_MODEL_NAME} from ${MODEL_URL}..."
-    mkdir -p "$MODEL_DIR"  # Ensure the directory exists before downloading
-    wget "$MODEL_URL" -P "$MODEL_DIR" || {
-        echo "Failed to download model from ${MODEL_URL}"
-        exit 1  # Exit if download fails
-    }
+    # Navigate to the directory where download_model.py is located
+    cd /home/user/app/app/utils
+
+    # Execute the download_model script
+    echo "Downloading the model..."
+    python download_model.py || { echo "Model download failed"; exit 1; }
+
+    # Navigate back to the app directory
+    cd /home/user/app
 else 
     echo "Model ${HF_MODEL_NAME} already present."
 fi