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

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchPypi,
   5   setuptools-scm,
   6   setuptools,
   7   pythonOlder,
   8   pydantic,
   9   openai,
  10   tiktoken,
  11   diskcache,
  12   cohere,
  13   google-auth,
  14   typer,
  15   pyyaml,
  16   transformers,
  17   fastapi,
  18   uvicorn,
  19   accelerate,
  20 }:
  21
  22 buildPythonPackage rec {
  23   pname = "llmx";
  24   version = "0.0.21a0";
  25   pyproject = true;
  26
  27   src = fetchPypi {
  28     inherit pname version;
  29     hash = "sha256-OEo6wIaDTktzAsP0rOmhxjFSHygTR/EpcRI6AXsu+6M=";
  30   };
  31
  32   disabled = pythonOlder "3.9";
  33
  34   build-system = [
  35     setuptools
  36     setuptools-scm
  37   ];
  38
  39   dependencies = [
  40     pydantic
  41     openai
  42     tiktoken
  43     diskcache
  44     cohere
  45     google-auth
  46     typer
  47     pyyaml
  48   ];
  49
  50   optional-dependencies = {
  51     web = [
  52       fastapi
  53       uvicorn
  54     ];
  55     transformers = [
  56       accelerate
  57       transformers
  58     ] ++ transformers.optional-dependencies.torch;
  59   };
  60
  61   # Tests of llmx try to access openai, google, etc.
  62   doCheck = false;
  63
  64   pythonImportsCheck = [ "llmx" ];
  65
  66   meta = with lib; {
  67     description = "Library for LLM Text Generation";
  68     homepage = "https://github.com/victordibia/llmx";
  69     mainProgram = "llmx";
  70     license = licenses.mit;
  71     maintainers = with maintainers; [ moraxyc ];
  72   };
  73 }