Skip to content

Commit 519cd0e

Browse files
mudlermudler
committed
feat(rerankers): Add new backend, support jina rerankers API
Signed-off-by: Ettore Di Giacinto <[email protected]>
1 parent d65214a commit 519cd0e

File tree

25 files changed

+630
-5
lines changed

25 files changed

+630
-5
lines changed

‎.github/workflows/test-extra.yml

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,37 @@ jobs:
7474
make --jobs=5 --output-sync=target -C backend/python/sentencetransformers
7575
make --jobs=5 --output-sync=target -C backend/python/sentencetransformers test
7676
77+
78+
tests-rerankers:
79+
runs-on: ubuntu-latest
80+
steps:
81+
- name: Clone
82+
uses: actions/checkout@v4
83+
with:
84+
submodules: true
85+
- name: Dependencies
86+
run: |
87+
sudo apt-get update
88+
sudo apt-get install build-essential ffmpeg
89+
curl https://repo.anaconda.com/pkgs/misc/gpgkeys/anaconda.asc | gpg --dearmor > conda.gpg && \
90+
sudo install -o root -g root -m 644 conda.gpg /usr/share/keyrings/conda-archive-keyring.gpg && \
91+
gpg --keyring /usr/share/keyrings/conda-archive-keyring.gpg --no-default-keyring --fingerprint 34161F5BF5EB1D4BFBBB8F0A8AEB4F8B29D82806 && \
92+
sudo /bin/bash -c 'echo "deb [arch=amd64 signed-by=/usr/share/keyrings/conda-archive-keyring.gpg] https://repo.anaconda.com/pkgs/misc/debrepo/conda stable main" > /etc/apt/sources.list.d/conda.list' && \
93+
sudo /bin/bash -c 'echo "deb [arch=amd64 signed-by=/usr/share/keyrings/conda-archive-keyring.gpg] https://repo.anaconda.com/pkgs/misc/debrepo/conda stable main" | tee -a /etc/apt/sources.list.d/conda.list' && \
94+
sudo apt-get update && \
95+
sudo apt-get install -y conda
96+
sudo apt-get install -y ca-certificates cmake curl patch python3-pip
97+
sudo apt-get install -y libopencv-dev
98+
pip install --user grpcio-tools
99+
100+
sudo rm -rfv /usr/bin/conda || true
101+
102+
- name: Test rerankers
103+
run: |
104+
export PATH=$PATH:/opt/conda/bin
105+
make --jobs=5 --output-sync=target -C backend/python/rerankers
106+
make --jobs=5 --output-sync=target -C backend/python/rerankers test
107+
77108
tests-diffusers:
78109
runs-on: ubuntu-latest
79110
steps:

‎Dockerfile

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ ARG TARGETVARIANT
1616

1717
ENV BUILD_TYPE=${BUILD_TYPE}
1818
ENV DEBIAN_FRONTEND=noninteractive
19-
ENV EXTERNAL_GRPC_BACKENDS="coqui:/build/backend/python/coqui/run.sh,huggingface-embeddings:/build/backend/python/sentencetransformers/run.sh,petals:/build/backend/python/petals/run.sh,transformers:/build/backend/python/transformers/run.sh,sentencetransformers:/build/backend/python/sentencetransformers/run.sh,autogptq:/build/backend/python/autogptq/run.sh,bark:/build/backend/python/bark/run.sh,diffusers:/build/backend/python/diffusers/run.sh,exllama:/build/backend/python/exllama/run.sh,vall-e-x:/build/backend/python/vall-e-x/run.sh,vllm:/build/backend/python/vllm/run.sh,mamba:/build/backend/python/mamba/run.sh,exllama2:/build/backend/python/exllama2/run.sh,transformers-musicgen:/build/backend/python/transformers-musicgen/run.sh,parler-tts:/build/backend/python/parler-tts/run.sh"
19+
ENV EXTERNAL_GRPC_BACKENDS="coqui:/build/backend/python/coqui/run.sh,huggingface-embeddings:/build/backend/python/sentencetransformers/run.sh,petals:/build/backend/python/petals/run.sh,transformers:/build/backend/python/transformers/run.sh,sentencetransformers:/build/backend/python/sentencetransformers/run.sh,rerankers:/build/backend/python/rerankers/run.sh,autogptq:/build/backend/python/autogptq/run.sh,bark:/build/backend/python/bark/run.sh,diffusers:/build/backend/python/diffusers/run.sh,exllama:/build/backend/python/exllama/run.sh,vall-e-x:/build/backend/python/vall-e-x/run.sh,vllm:/build/backend/python/vllm/run.sh,mamba:/build/backend/python/mamba/run.sh,exllama2:/build/backend/python/exllama2/run.sh,transformers-musicgen:/build/backend/python/transformers-musicgen/run.sh,parler-tts:/build/backend/python/parler-tts/run.sh"
2020

2121
ARG GO_TAGS="stablediffusion tinydream tts"
2222

@@ -259,6 +259,9 @@ RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
259259
RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
260260
make -C backend/python/sentencetransformers \
261261
; fi
262+
RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
263+
make -C backend/python/rerankers \
264+
; fi
262265
RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
263266
make -C backend/python/transformers \
264267
; fi

‎Makefile

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -437,10 +437,10 @@ protogen-go-clean:
437437
$(RM) bin/*
438438

439439
.PHONY: protogen-python
440-
protogen-python: autogptq-protogen bark-protogen coqui-protogen diffusers-protogen exllama-protogen exllama2-protogen mamba-protogen petals-protogen sentencetransformers-protogen transformers-protogen parler-tts-protogen transformers-musicgen-protogen vall-e-x-protogen vllm-protogen
440+
protogen-python: autogptq-protogen bark-protogen coqui-protogen diffusers-protogen exllama-protogen exllama2-protogen mamba-protogen petals-protogen rerankers-protogen sentencetransformers-protogen transformers-protogen parler-tts-protogen transformers-musicgen-protogen vall-e-x-protogen vllm-protogen
441441

442442
.PHONY: protogen-python-clean
443-
protogen-python-clean: autogptq-protogen-clean bark-protogen-clean coqui-protogen-clean diffusers-protogen-clean exllama-protogen-clean exllama2-protogen-clean mamba-protogen-clean petals-protogen-clean sentencetransformers-protogen-clean transformers-protogen-clean transformers-musicgen-protogen-clean parler-tts-protogen-clean vall-e-x-protogen-clean vllm-protogen-clean
443+
protogen-python-clean: autogptq-protogen-clean bark-protogen-clean coqui-protogen-clean diffusers-protogen-clean exllama-protogen-clean exllama2-protogen-clean mamba-protogen-clean petals-protogen-clean sentencetransformers-protogen-clean rerankers-protogen-clean transformers-protogen-clean transformers-musicgen-protogen-clean parler-tts-protogen-clean vall-e-x-protogen-clean vllm-protogen-clean
444444

445445
.PHONY: autogptq-protogen
446446
autogptq-protogen:
@@ -506,6 +506,14 @@ petals-protogen:
506506
petals-protogen-clean:
507507
$(MAKE) -C backend/python/petals protogen-clean
508508

509+
.PHONY: rerankers-protogen
510+
rerankers-protogen:
511+
$(MAKE) -C backend/python/rerankers protogen
512+
513+
.PHONY: rerankers-protogen-clean
514+
rerankers-protogen-clean:
515+
$(MAKE) -C backend/python/rerankers protogen-clean
516+
509517
.PHONY: sentencetransformers-protogen
510518
sentencetransformers-protogen:
511519
$(MAKE) -C backend/python/sentencetransformers protogen
@@ -564,6 +572,7 @@ prepare-extra-conda-environments: protogen-python
564572
$(MAKE) -C backend/python/vllm
565573
$(MAKE) -C backend/python/mamba
566574
$(MAKE) -C backend/python/sentencetransformers
575+
$(MAKE) -C backend/python/rerankers
567576
$(MAKE) -C backend/python/transformers
568577
$(MAKE) -C backend/python/transformers-musicgen
569578
$(MAKE) -C backend/python/parler-tts

‎aio/cpu/rerank.yaml

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
name: jina-reranker-v1-base-en
2+
backend: rerankers
3+
parameters:
4+
model: cross-encoder
5+
6+
usage: |
7+
You can test this model with curl like this:
8+
9+
curl http://localhost:8080/v1/rerank \
10+
-H "Content-Type: application/json" \
11+
-d '{
12+
"model": "jina-reranker-v1-base-en",
13+
"query": "Organic skincare products for sensitive skin",
14+
"documents": [
15+
"Eco-friendly kitchenware for modern homes",
16+
"Biodegradable cleaning supplies for eco-conscious consumers",
17+
"Organic cotton baby clothes for sensitive skin",
18+
"Natural organic skincare range for sensitive skin",
19+
"Tech gadgets for smart homes: 2024 edition",
20+
"Sustainable gardening tools and compost solutions",
21+
"Sensitive skin-friendly facial cleansers and toners",
22+
"Organic food wraps and storage solutions",
23+
"All-natural pet food for dogs with allergies",
24+
"Yoga mats made from recycled materials"
25+
],
26+
"top_n": 3
27+
}'

‎aio/entrypoint.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ detect_gpu
129129
detect_gpu_size
130130

131131
PROFILE="${PROFILE:-$GPU_SIZE}" # default to cpu
132-
export MODELS="${MODELS:-/aio/${PROFILE}/embeddings.yaml,/aio/${PROFILE}/text-to-speech.yaml,/aio/${PROFILE}/image-gen.yaml,/aio/${PROFILE}/text-to-text.yaml,/aio/${PROFILE}/speech-to-text.yaml,/aio/${PROFILE}/vision.yaml}"
132+
export MODELS="${MODELS:-/aio/${PROFILE}/embeddings.yaml,/aio/${PROFILE}/rerank.yaml,/aio/${PROFILE}/text-to-speech.yaml,/aio/${PROFILE}/image-gen.yaml,/aio/${PROFILE}/text-to-text.yaml,/aio/${PROFILE}/speech-to-text.yaml,/aio/${PROFILE}/vision.yaml}"
133133

134134
check_vars
135135

‎aio/gpu-8g/rerank.yaml

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
name: jina-reranker-v1-base-en
2+
backend: rerankers
3+
parameters:
4+
model: cross-encoder
5+
6+
usage: |
7+
You can test this model with curl like this:
8+
9+
curl http://localhost:8080/v1/rerank \
10+
-H "Content-Type: application/json" \
11+
-d '{
12+
"model": "jina-reranker-v1-base-en",
13+
"query": "Organic skincare products for sensitive skin",
14+
"documents": [
15+
"Eco-friendly kitchenware for modern homes",
16+
"Biodegradable cleaning supplies for eco-conscious consumers",
17+
"Organic cotton baby clothes for sensitive skin",
18+
"Natural organic skincare range for sensitive skin",
19+
"Tech gadgets for smart homes: 2024 edition",
20+
"Sustainable gardening tools and compost solutions",
21+
"Sensitive skin-friendly facial cleansers and toners",
22+
"Organic food wraps and storage solutions",
23+
"All-natural pet food for dogs with allergies",
24+
"Yoga mats made from recycled materials"
25+
],
26+
"top_n": 3
27+
}'

‎aio/intel/rerank.yaml

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
name: jina-reranker-v1-base-en
2+
backend: rerankers
3+
parameters:
4+
model: cross-encoder
5+
6+
usage: |
7+
You can test this model with curl like this:
8+
9+
curl http://localhost:8080/v1/rerank \
10+
-H "Content-Type: application/json" \
11+
-d '{
12+
"model": "jina-reranker-v1-base-en",
13+
"query": "Organic skincare products for sensitive skin",
14+
"documents": [
15+
"Eco-friendly kitchenware for modern homes",
16+
"Biodegradable cleaning supplies for eco-conscious consumers",
17+
"Organic cotton baby clothes for sensitive skin",
18+
"Natural organic skincare range for sensitive skin",
19+
"Tech gadgets for smart homes: 2024 edition",
20+
"Sustainable gardening tools and compost solutions",
21+
"Sensitive skin-friendly facial cleansers and toners",
22+
"Organic food wraps and storage solutions",
23+
"All-natural pet food for dogs with allergies",
24+
"Yoga mats made from recycled materials"
25+
],
26+
"top_n": 3
27+
}'

‎backend/backend.proto

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,30 @@ service Backend {
2323
rpc StoresDelete(StoresDeleteOptions) returns (Result) {}
2424
rpc StoresGet(StoresGetOptions) returns (StoresGetResult) {}
2525
rpc StoresFind(StoresFindOptions) returns (StoresFindResult) {}
26+
27+
rpc Rerank(RerankRequest) returns (RerankResult) {}
28+
}
29+
30+
message RerankRequest {
31+
string query = 1;
32+
repeated string documents = 2;
33+
int32 top_n = 3;
34+
}
35+
36+
message RerankResult {
37+
Usage usage = 1;
38+
repeated DocumentResult results = 2;
39+
}
40+
41+
message Usage {
42+
int32 total_tokens = 1;
43+
int32 prompt_tokens = 2;
44+
}
45+
46+
message DocumentResult {
47+
int32 index = 1;
48+
string text = 2;
49+
float relevance_score = 3;
2650
}
2751

2852
message StoresKey {

‎backend/python/common-env/transformers/transformers-nvidia.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,4 +120,6 @@ dependencies:
120120
- transformers>=4.38.2 # Updated Version
121121
- transformers_stream_generator==0.0.5
122122
- xformers==0.0.23.post1
123+
- rerankers[transformers]
124+
- pydantic
123125
prefix: /opt/conda/envs/transformers

‎backend/python/common-env/transformers/transformers-rocm.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,4 +108,6 @@ dependencies:
108108
- transformers>=4.38.2 # Updated Version
109109
- transformers_stream_generator==0.0.5
110110
- xformers==0.0.23.post1
111+
- rerankers[transformers]
112+
- pydantic
111113
prefix: /opt/conda/envs/transformers

0 commit comments

Comments
 (0)