pkgs/development/python-modules/sentence-transformers/default.nix

   1 {
   2   lib,
   3   stdenv,
   4   buildPythonPackage,
   5   fetchFromGitHub,
   6
   7   # build-system
   8   setuptools,
   9
  10   # dependencies
  11   huggingface-hub,
  12   nltk,
  13   numpy,
  14   scikit-learn,
  15   scipy,
  16   sentencepiece,
  17   tokenizers,
  18   torch,
  19   tqdm,
  20   transformers,
  21
  22   # tests
  23   accelerate,
  24   datasets,
  25   pytestCheckHook,
  26   pytest-cov,
  27 }:
  28
  29 buildPythonPackage rec {
  30   pname = "sentence-transformers";
  31   version = "3.2.0";
  32   pyproject = true;
  33
  34   src = fetchFromGitHub {
  35     owner = "UKPLab";
  36     repo = "sentence-transformers";
  37     rev = "refs/tags/v${version}";
  38     hash = "sha256-T9y+usEOqz53Wv5gObD5o3gAlABHM3apoq/ySD5ynVg=";
  39   };
  40
  41   build-system = [ setuptools ];
  42
  43   dependencies = [
  44     huggingface-hub
  45     nltk
  46     numpy
  47     scikit-learn
  48     scipy
  49     sentencepiece
  50     tokenizers
  51     torch
  52     tqdm
  53     transformers
  54   ];
  55
  56   nativeCheckInputs = [
  57     accelerate
  58     datasets
  59     pytestCheckHook
  60     pytest-cov
  61   ];
  62
  63   pythonImportsCheck = [ "sentence_transformers" ];
  64
  65   disabledTests = [
  66     # Tests require network access
  67     "test_cmnrl_same_grad"
  68     "test_LabelAccuracyEvaluator"
  69     "test_model_card_reuse"
  70     "test_paraphrase_mining"
  71     "test_ParaphraseMiningEvaluator"
  72     "test_simple_encode"
  73     "test_trainer"
  74     "test_trainer_invalid_column_names"
  75     "test_trainer_multi_dataset_errors"
  76   ];
  77
  78   disabledTestPaths = [
  79     # Tests require network access
  80     "tests/evaluation/test_information_retrieval_evaluator.py"
  81     "tests/test_compute_embeddings.py"
  82     "tests/test_cross_encoder.py"
  83     "tests/test_model_card_data.py"
  84     "tests/test_multi_process.py"
  85     "tests/test_pretrained_stsb.py"
  86     "tests/test_sentence_transformer.py"
  87     "tests/test_train_stsb.py"
  88   ];
  89
  90   # Sentence-transformer needs a writable hf_home cache
  91   postInstall = ''
  92     export HF_HOME=$(mktemp -d)
  93   '';
  94
  95   meta = {
  96     description = "Multilingual Sentence & Image Embeddings with BERT";
  97     homepage = "https://github.com/UKPLab/sentence-transformers";
  98     changelog = "https://github.com/UKPLab/sentence-transformers/releases/tag/${version}";
  99     license = lib.licenses.asl20;
 100     maintainers = with lib.maintainers; [ dit7ya ];
 101     # Segmentation fault at import
 102     broken = stdenv.hostPlatform.system == "x86_64-darwin";
 103   };
 104 }