1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 setuptools,
8
9 # dependencies
10 accelerate,
11 datasets,
12 huggingface-hub,
13 optimum,
14 pillow,
15 scikit-learn,
16 scipy,
17 torch,
18 tqdm,
19 transformers,
20 typing-extensions,
21
22 # tests
23 pytestCheckHook,
24 pytest-cov-stub,
25}:
26
27buildPythonPackage rec {
28 pname = "sentence-transformers";
29 version = "4.1.0";
30 pyproject = true;
31
32 src = fetchFromGitHub {
33 owner = "UKPLab";
34 repo = "sentence-transformers";
35 tag = "v${version}";
36 hash = "sha256-9Mg3+7Yxf195h4cUNLP/Z1PrauxanHJfS8OV2JIwRj4=";
37 };
38
39 build-system = [ setuptools ];
40
41 dependencies = [
42 huggingface-hub
43 pillow
44 scikit-learn
45 scipy
46 torch
47 tqdm
48 transformers
49 typing-extensions
50 ];
51
52 optional-dependencies = {
53 train = [
54 accelerate
55 datasets
56 ];
57 onnx = [ optimum ] ++ optimum.optional-dependencies.onnxruntime;
58 # onnx-gpu = [ optimum ] ++ optimum.optional-dependencies.onnxruntime-gpu;
59 # openvino = [ optimum-intel ] ++ optimum-intel.optional-dependencies.openvino;
60 };
61
62 nativeCheckInputs = [
63 pytest-cov-stub
64 pytestCheckHook
65 ] ++ lib.flatten (builtins.attrValues optional-dependencies);
66
67 pythonImportsCheck = [ "sentence_transformers" ];
68
69 disabledTests = [
70 # Tests require network access
71 "test_LabelAccuracyEvaluator"
72 "test_ParaphraseMiningEvaluator"
73 "test_TripletEvaluator"
74 "test_cmnrl_same_grad"
75 "test_forward"
76 "test_initialization_with_embedding_dim"
77 "test_initialization_with_embedding_weights"
78 "test_loading_model2vec"
79 "test_model_card_base"
80 "test_model_card_reuse"
81 "test_nanobeir_evaluator"
82 "test_paraphrase_mining"
83 "test_pretrained_model"
84 "test_save_and_load"
85 "test_simple_encode"
86 "test_tokenize"
87 "test_trainer"
88 "test_trainer_invalid_column_names"
89 "test_trainer_multi_dataset_errors"
90 ];
91
92 disabledTestPaths = [
93 # Tests require network access
94 "tests/cross_encoder/test_cross_encoder.py"
95 "tests/cross_encoder/test_train_stsb.py"
96 "tests/evaluation/test_information_retrieval_evaluator.py"
97 "tests/test_compute_embeddings.py"
98 "tests/test_model_card_data.py"
99 "tests/test_multi_process.py"
100 "tests/test_pretrained_stsb.py"
101 "tests/test_sentence_transformer.py"
102 "tests/test_train_stsb.py"
103 ];
104
105 # Sentence-transformer needs a writable hf_home cache
106 postInstall = ''
107 export HF_HOME=$(mktemp -d)
108 '';
109
110 meta = {
111 description = "Multilingual Sentence & Image Embeddings with BERT";
112 homepage = "https://github.com/UKPLab/sentence-transformers";
113 changelog = "https://github.com/UKPLab/sentence-transformers/releases/tag/${src.tag}";
114 license = lib.licenses.asl20;
115 maintainers = with lib.maintainers; [ dit7ya ];
116 };
117}