tjh.dev
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 = "5.1.0";
30 pyproject = true;
31
32 src = fetchFromGitHub {
33 owner = "UKPLab";
34 repo = "sentence-transformers";
35 tag = "v${version}";
36 hash = "sha256-snowpTdHelcFjo1+hvqpoVt5ROB0f91yt0GsIvA5cso=";
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 ]
66 ++ lib.flatten (builtins.attrValues optional-dependencies);
67
68 pythonImportsCheck = [ "sentence_transformers" ];
69
70 disabledTests = [
71 # Tests require network access
72 "test_LabelAccuracyEvaluator"
73 "test_ParaphraseMiningEvaluator"
74 "test_TripletEvaluator"
75 "test_cmnrl_same_grad"
76 "test_forward"
77 "test_initialization_with_embedding_dim"
78 "test_initialization_with_embedding_weights"
79 "test_loading_model2vec"
80 "test_mine_hard_negatives_with_prompt"
81 "test_model_card_base"
82 "test_model_card_reuse"
83 "test_nanobeir_evaluator"
84 "test_paraphrase_mining"
85 "test_pretrained_model"
86 "test_router_as_middle_module"
87 "test_router_backwards_compatibility"
88 "test_router_encode"
89 "test_router_load_with_config"
90 "test_router_save_load"
91 "test_router_save_load_with_custom_default_route"
92 "test_router_save_load_with_multiple_modules_per_route"
93 "test_router_save_load_without_default_route"
94 "test_router_with_trainer"
95 "test_router_with_trainer_without_router_mapping"
96 "test_save_and_load"
97 "test_simple_encode"
98 "test_tokenize"
99 "test_train_stsb"
100 "test_trainer"
101 "test_trainer_invalid_column_names"
102 "test_trainer_multi_dataset_errors"
103
104 # Assertion error: Sparse operations take too long
105 # (namely, load-sensitive test)
106 "test_performance_with_large_vectors"
107 ];
108
109 disabledTestPaths = [
110 # Tests require network access
111 "tests/cross_encoder/test_cross_encoder.py"
112 "tests/cross_encoder/test_train_stsb.py"
113 "tests/evaluation/test_information_retrieval_evaluator.py"
114 "tests/sparse_encoder/models/test_csr.py"
115 "tests/sparse_encoder/models/test_sparse_static_embedding.py"
116 "tests/sparse_encoder/test_opensearch_models.py"
117 "tests/sparse_encoder/test_pretrained.py"
118 "tests/sparse_encoder/test_sparse_encoder.py"
119 "tests/test_compute_embeddings.py"
120 "tests/test_model_card_data.py"
121 "tests/test_multi_process.py"
122 "tests/test_pretrained_stsb.py"
123 "tests/test_sentence_transformer.py"
124 "tests/test_train_stsb.py"
125 "tests/util/test_hard_negatives.py"
126 ];
127
128 # Sentence-transformer needs a writable hf_home cache
129 postInstall = ''
130 export HF_HOME=$(mktemp -d)
131 '';
132
133 meta = {
134 description = "Multilingual Sentence & Image Embeddings with BERT";
135 homepage = "https://github.com/UKPLab/sentence-transformers";
136 changelog = "https://github.com/UKPLab/sentence-transformers/releases/tag/${src.tag}";
137 license = lib.licenses.asl20;
138 maintainers = with lib.maintainers; [ dit7ya ];
139 };
140}