depot/third_party/nixpkgs/pkgs/development/python-modules/openllm/default.nix

{ lib
, buildPythonPackage
, hatch-fancy-pypi-readme
, hatch-vcs
, hatchling
, pytestCheckHook
, pythonOlder
, pythonRelaxDepsHook
, accelerate
, bentoml
, bitsandbytes
, build
, click
, ctranslate2
, datasets
, docker
, einops
, ghapi
, huggingface-hub
, hypothesis
, ipython
, jupyter
, jupytext
, nbformat
, notebook
, openai
, openllm-client
, openllm-core
, optimum
, peft
, pytest-mock
, pytest-randomly
, pytest-rerunfailures
, pytest-xdist
, safetensors
, scipy
, sentencepiece
, soundfile
, syrupy
, tabulate
, tiktoken
, transformers
, openai-triton
, xformers
}:

buildPythonPackage rec {
  inherit (openllm-core) src version;
  pname = "openllm";
  pyproject = true;

  disabled = pythonOlder "3.8";

  sourceRoot = "source/openllm-python";

  nativeBuildInputs = [
    hatch-fancy-pypi-readme
    hatch-vcs
    hatchling
    pythonRelaxDepsHook
  ];

  pythonRemoveDeps = [
    # remove cuda-python as it has an unfree license
    "cuda-python"
  ];

  propagatedBuildInputs = [
    accelerate
    bentoml
    bitsandbytes
    build
    click
    einops
    ghapi
    openllm-client
    openllm-core
    optimum
    safetensors
    scipy
    sentencepiece
    tabulate
    transformers
  ] ++ bentoml.optional-dependencies.io
  ++ tabulate.optional-dependencies.widechars
  ++ transformers.optional-dependencies.tokenizers
  ++ transformers.optional-dependencies.torch;

  passthru.optional-dependencies = {
    agents = [
      # diffusers
      soundfile
      transformers
    ] ++ transformers.optional-dependencies.agents;
    awq = [
      # autoawq
    ];
    baichuan = [
      # cpm-kernels
    ];
    chatglm = [
      # cpm-kernels
    ];
    ctranslate = [
      ctranslate2
    ];
    falcon = [
      xformers
    ];
    fine-tune = [
      datasets
      huggingface-hub
      peft
      # trl
    ];
    ggml = [
      # ctransformers
    ];
    gptq = [
      # auto-gptq
    ]; # ++ autogptq.optional-dependencies.triton;
    grpc = [
      bentoml
    ] ++ bentoml.optional-dependencies.grpc;
    mpt = [
      openai-triton
    ];
    openai = [
      openai
      tiktoken
    ] ++ openai.optional-dependencies.datalib;
    playground = [
      ipython
      jupyter
      jupytext
      nbformat
      notebook
    ];
    starcoder = [
      bitsandbytes
    ];
    vllm = [
      # vllm
    ];
    full = with passthru.optional-dependencies; (
      agents ++ awq ++ baichuan ++ chatglm ++ ctranslate ++ falcon ++ fine-tune ++ ggml ++ gptq ++ mpt ++ openai ++ playground ++ starcoder ++ vllm
    );
    all = passthru.optional-dependencies.full;
  };

  nativeCheckInputs = [
    docker
    hypothesis
    pytest-mock
    pytest-randomly
    pytest-rerunfailures
    pytest-xdist
    pytestCheckHook
    syrupy
  ];

  preCheck = ''
    export HOME=$TMPDIR
    # skip GPUs test on CI
    export GITHUB_ACTIONS=1
    # disable hypothesis' deadline
    export CI=1
  '';

  disabledTestPaths = [
    # require network access
    "tests/models"
  ];

  disabledTests = [
    # incompatible with recent TypedDict
    # https://github.com/bentoml/OpenLLM/blob/f3fd32d596253ae34c68e2e9655f19f40e05f666/openllm-python/tests/configuration_test.py#L18-L21
    "test_missing_default"
  ];

  pythonImportsCheck = [ "openllm" ];

  meta = with lib; {
    description = "Operating LLMs in production";
    homepage = "https://github.com/bentoml/OpenLLM/tree/main/openllm-python";
    changelog = "https://github.com/bentoml/OpenLLM/blob/${src.rev}/CHANGELOG.md";
    license = licenses.asl20;
    maintainers = with maintainers; [ happysalada natsukium ];
  };
}