pkgs/development/python-modules/argilla/default.nix

   1 { lib
   2 , buildPythonPackage
   3 , fetchFromGitHub
   4 , pythonOlder
   5 , pythonRelaxDepsHook
   6 , deprecated
   7 , rich
   8 , backoff
   9 , packaging
  10 , pydantic
  11 , typer
  12 , tqdm
  13 , wrapt
  14 , numpy
  15 , httpx
  16 , pandas
  17 , monotonic
  18 # optional-dependencies
  19 , fastapi
  20 , opensearch-py
  21 , elasticsearch8
  22 , uvicorn
  23 , smart-open
  24 , brotli-asgi
  25 , alembic
  26 , sqlalchemy
  27 , greenlet
  28 , aiosqlite
  29 , luqum
  30 , scikit-learn
  31 , aiofiles
  32 , pyyaml
  33 , python-multipart
  34 , python-jose
  35 , passlib
  36 , psutil
  37 # , segment-analytics-python
  38 , asyncpg
  39 , psycopg2
  40 , schedule
  41 , prodict
  42 , cleanlab
  43 , datasets
  44 , huggingface-hub
  45 # , flair
  46 , faiss
  47 , flyingsquid
  48 , pgmpy
  49 , plotly
  50 , snorkel
  51 , spacy
  52 , transformers
  53 , evaluate
  54 , seqeval
  55 # , setfit
  56 # , span_marker
  57 , openai
  58 , peft
  59 # test dependencies
  60 , pytestCheckHook
  61 , pytest-cov
  62 , pytest-mock
  63 , pytest-asyncio
  64 , factory-boy
  65 }:
  66 let
  67   pname = "argilla";
  68   version = "1.19.0";
  69   optional-dependencies = {
  70     server = [
  71       fastapi
  72       opensearch-py
  73       elasticsearch8
  74       uvicorn
  75       smart-open
  76       brotli-asgi
  77       alembic
  78       sqlalchemy
  79       greenlet
  80       aiosqlite
  81       luqum
  82       scikit-learn
  83       aiofiles
  84       pyyaml
  85       python-multipart
  86       python-jose
  87       passlib
  88       psutil
  89       # segment-analytics-python
  90     ] ++
  91       elasticsearch8.optional-dependencies.async ++
  92       uvicorn.optional-dependencies.standard ++
  93       python-jose.optional-dependencies.cryptography ++
  94       passlib.optional-dependencies.bcrypt;
  95     postgresql = [ asyncpg psycopg2 ];
  96     listeners = [ schedule prodict ];
  97     integrations = [
  98       pyyaml
  99       cleanlab
 100       datasets
 101       huggingface-hub
 102       # flair
 103       faiss
 104       flyingsquid
 105       pgmpy
 106       plotly
 107       snorkel
 108       spacy
 109       transformers
 110       evaluate
 111       seqeval
 112       # setfit
 113       # span_marker
 114       openai
 115       peft
 116     ] ++ transformers.optional-dependencies.torch;
 117   };
 118 in
 119 buildPythonPackage {
 120   inherit pname version;
 121   format = "setuptools";
 122
 123   disabled = pythonOlder "3.8";
 124
 125   src = fetchFromGitHub {
 126     owner = "argilla-io";
 127     repo = pname;
 128     rev = "refs/tags/v${version}";
 129     hash = "sha256-Idl5Tm1XWgBLVgHPbXiyt9MW4J5wZdPb2J7iIDBnorg=";
 130   };
 131
 132   pythonRelaxDeps = [
 133     "typer"
 134     "rich"
 135     "numpy"
 136   ];
 137
 138   nativeBuildInputs = [
 139     pythonRelaxDepsHook
 140   ];
 141
 142   propagatedBuildInputs = [
 143     httpx
 144     deprecated
 145     packaging
 146     pandas
 147     pydantic
 148     wrapt
 149     numpy
 150     tqdm
 151     backoff
 152     monotonic
 153     rich
 154     typer
 155   ];
 156
 157   # still quite a bit of optional dependencies missing
 158   doCheck = false;
 159
 160   preCheck = ''
 161     export HOME=$(mktemp -d)
 162   '';
 163
 164   nativeCheckInputs = [
 165     pytestCheckHook
 166     pytest-cov
 167     pytest-mock
 168     pytest-asyncio
 169     factory-boy
 170   ]
 171     ++ optional-dependencies.server
 172     ++ optional-dependencies.postgresql
 173     ++ optional-dependencies.listeners
 174     ++ optional-dependencies.integrations;
 175
 176   pytestFlagsArray = [ "--ignore=tests/server/datasets/test_dao.py" ];
 177
 178   passthru.optional-dependencies = optional-dependencies;
 179
 180   meta = with lib; {
 181     description = "Argilla: the open-source data curation platform for LLMs";
 182     homepage = "https://github.com/argilla-io/argilla";
 183     changelog = "https://github.com/argilla-io/argilla/releases/tag/v${version}";
 184     license = licenses.asl20;
 185     maintainers = with maintainers; [ happysalada ];
 186   };
 187 }