Build a custom Pod template

# Run application after services start
COPY run.sh /app/run.sh
RUN chmod +x /app/run.sh
CMD ["/app/run.sh"]

#!/bin/bash
# Start base image services (Jupyter/SSH) in background
/start.sh &

# Wait for services to start
sleep 2

# Run your application
python /app/main.py

# Wait for background processes
wait

# Clear entrypoint and run application only
ENTRYPOINT []
CMD ["python", "/app/main.py"]

# Set Hugging Face cache directory
ENV HF_HOME=/app/models
ENV HF_HUB_ENABLE_HF_TRANSFER=0

# Pre-download model during build
RUN python -c "from transformers import pipeline; pipeline('sentiment-analysis', model='distilbert-base-uncased-finetuned-sst-2-english')"

# Create model directory and download files
RUN mkdir -p /app/models/distilbert-model && \
    cd /app/models/distilbert-model && \
    wget -q https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json && \
    wget -q https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/model.safetensors && \
    wget -q https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/tokenizer_config.json && \
    wget -q https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/vocab.txt

"""
Example Pod template application with sentiment analysis.
"""

import sys
import torch
import time
import signal
from transformers import pipeline

def main():
    print("Hello from your custom Runpod template!")
    print(f"Python version: {sys.version.split()[0]}")
    print(f"PyTorch version: {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"CUDA version: {torch.version.cuda}")
        print(f"GPU device: {torch.cuda.get_device_name(0)}")
    
    # Initialize model
    print("\nLoading sentiment analysis model...")
    device = 0 if torch.cuda.is_available() else -1
    
    # MODEL LOADING OPTIONS:
    
    # OPTION 1: From Hugging Face Hub cache (default)
    # Bakes the model into the container image using transformers pipeline
    # Behavior: Loads model from the cache, requires local_files_only=True
    classifier = pipeline(
        "sentiment-analysis",
        model="distilbert-base-uncased-finetuned-sst-2-english",
        device=device,
        model_kwargs={"local_files_only": True},
    )

    # OPTION 2: From a local directory
    # Download the model files using wget, loads them from the local directory
    # Behavior: Loads directly from /app/models/distilbert-model
    # To use: Uncomment the pipeline object below, comment OPTION 1 above
    # classifier = pipeline('sentiment-analysis',
    #                       model='/app/models/distilbert-model',
    #                       device=device)
    
    print("Model loaded successfully!")
    
    # Example inference
    test_texts = [
        "This is a wonderful experience!",
        "I really don't like this at all.",
        "The weather is nice today.",
    ]
    
    print("\n--- Running sentiment analysis ---")
    for text in test_texts:
        result = classifier(text)
        print(f"Text: {text}")
        print(f"Result: {result[0]['label']} (confidence: {result[0]['score']:.4f})\n")
    
    print("Container is running. Press Ctrl+C to stop.")
    
    # Keep container running
    def signal_handler(sig, frame):
        print("\nShutting down...")
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try:
        while True:
            time.sleep(60)
    except KeyboardInterrupt:
        signal_handler(None, None)

if __name__ == "__main__":
    main()