cog.yaml

# Configuration for Cog ⚙️
# Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md

build:
  gpu: true

  python_version: "3.12"

  python_packages:
    - "torch>=2.0.1"
    - "accelerate>=0.21.0"
    - "bitsandbytes>=0.41.0"
    - "deepspeed>=0.9.5"
    - "einops-exts>=0.0.4"
    - "einops>=0.6.1"
    - "gradio>=3.35.2"
    - "gradio_client>=0.2.9"
    - "httpx>=0.24.0"
    - "markdown2>=2.4.10"
    - "numpy>=1.26.0"
    - "peft>=0.4.0"
    - "scikit-learn>=1.2.2"
    - "sentencepiece>=0.1.99"
    - "shortuuid>=1.0.11"
    - "timm>=0.6.13"
    - "tokenizers>=0.13.3"
    - "torch>=2.0.1"
    - "torchvision>=0.15.2"
    - "transformers>=4.31.0"
    - "wandb>=0.15.12"
    - "wavedrom>=2.0.3.post3"
    - "Pygments>=2.16.1"
  run:
    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.0.3/pget" && chmod +x /usr/local/bin/pget

# predict.py defines how predictions are run on your model
predict: "predict.py:Predictor"