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

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