Dockerfile.d/gemma_check.py

import os, subprocess

# Check if gemma-2b already exists
gemma_dir = os.path.join(os.getcwd(), ".hf_home/google/gemma-2b")
print("Checking", gemma_dir, "...")
if os.path.isdir(gemma_dir):
    print("... already exists")
    exit()

# Check if in a Hunggingface Space
if os.path.exists("/run/secrets/HUGGINGFACE_TOKEN"):
    print("... prefetch not needed")
    exit()

# Check if /run/secrets/dotenv file exists
if not os.path.isfile("/run/secrets/dotenv"):
    print("... can't prefetch, can't find --secret dotenv file")
    exit(1)

# Read the dotenv file and export the variables
with open("/run/secrets/dotenv") as dotenv_file:
    for line in dotenv_file:
        if '=' in line:
            key, value = line.split("=", 1)
            os.environ[key.strip()] = value.split("#", 1)[0].strip()
    subprocess.run(["python", "download-huggingface-model.py", "google/gemma-2b"])