|
--- |
|
license: apache-2.0 |
|
language: |
|
- en |
|
pipeline_tag: feature-extraction |
|
tags: |
|
- code |
|
--- |
|
|
|
<br><br> |
|
|
|
<p align="center"> |
|
<img src="https://github.com/jina-ai/finetuner/blob/main/docs/_static/finetuner-logo-ani.svg?raw=true" alt="Finetuner logo: Finetuner helps you to create experiments in order to improve embeddings on search tasks. It accompanies you to deliver the last mile of performance-tuning for neural search applications." width="150px"> |
|
</p> |
|
|
|
|
|
<p align="center"> |
|
<b>The text embedding set trained by <a href="https://jina.ai/"><b>Jina AI</b></a>, <a href="https://github.com/jina-ai/finetuner"><b>Finetuner</b></a> team.</b> |
|
</p> |
|
|
|
|
|
## Intended Usage & Model Info |
|
|
|
`jina-embeddings-v2-base-code` is an multilingual **embedding model** speaks English and 29 widely used programming languages supporting **8192 sequence length**. |
|
It is based on a Bert architecture (JinaBert) that supports the symmetric bidirectional variant of [ALiBi](https://arxiv.org/abs/2108.12409) to allow longer sequence length. |
|
The backbone `jina-bert-v2-base-code` is pretrained on the [github-code](https://huggingface.co/datasets/codeparrot/github-code) dataset. |
|
The model is further trained on Jina AI's collection of more than 150 millions of coding question answer and docstring source code pairs. |
|
These pairs were obtained from various domains and were carefully selected through a thorough cleaning process. |
|
|
|
The embedding model was trained using 512 sequence length, but extrapolates to 8k sequence length (or even longer) thanks to ALiBi. |
|
This makes our model useful for a range of use cases, especially when processing long documents is needed, including technical question answering and code search. |
|
|
|
This model has 137 million parameters, which enables fast and memory efficient inference, while delivering impressive performance. |
|
Additionally, we provide the following embedding models: |
|
|
|
**V1 (Based on T5, 512 Seq)** |
|
|
|
- [`jina-embeddings-v1-small-en`](https://huggingface.co/jinaai/jina-embedding-s-en-v1): 35 million parameters. |
|
- [`jina-embeddings-v1-base-en`](https://huggingface.co/jinaai/jina-embedding-b-en-v1): 110 million parameters. |
|
- [`jina-embeddings-v1-large-en`](https://huggingface.co/jinaai/jina-embedding-l-en-v1): 330 million parameters. |
|
|
|
**V2 (Based on JinaBert, 8k Seq)** |
|
|
|
- [`jina-embeddings-v2-small-en`](https://huggingface.co/jinaai/jina-embeddings-v2-small-en): 33 million parameters. |
|
- [`jina-embeddings-v2-base-en`](https://huggingface.co/jinaai/jina-embeddings-v2-base-en): 137 million parameters. |
|
- [`jina-embeddings-v2-base-code`](https://huggingface.co/jinaai/jina-embeddings-v2-base-code): 137 million parameters **(you are here)**. |
|
|
|
## Data & Parameters |
|
|
|
Jina Embeddings V2 [technical report](https://arxiv.org/abs/2310.19923) |
|
|
|
## Usage |
|
|
|
**<details><summary>Please apply mean pooling when integrating the model.</summary>** |
|
<p> |
|
|
|
### Why mean pooling? |
|
|
|
`mean poooling` takes all token embeddings from model output and averaging them at sentence/paragraph level. |
|
It has been proved to be the most effective way to produce high-quality sentence embeddings. |
|
We offer an `encode` function to deal with this. |
|
|
|
However, if you would like to do it without using the default `encode` function: |
|
|
|
```python |
|
import torch |
|
import torch.nn.functional as F |
|
from transformers import AutoTokenizer, AutoModel |
|
|
|
def mean_pooling(model_output, attention_mask): |
|
token_embeddings = model_output[0] |
|
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() |
|
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) |
|
|
|
sentences = ['How is the weather today?', 'What is the current weather like today?'] |
|
|
|
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-small-en') |
|
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-small-en', trust_remote_code=True) |
|
|
|
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt') |
|
|
|
with torch.no_grad(): |
|
model_output = model(**encoded_input) |
|
|
|
embeddings = mean_pooling(model_output, encoded_input['attention_mask']) |
|
embeddings = F.normalize(embeddings, p=2, dim=1) |
|
``` |
|
|
|
</p> |
|
</details> |
|
|
|
You can use Jina Embedding models directly from transformers package: |
|
```python |
|
!pip install transformers |
|
from transformers import AutoModel |
|
from numpy.linalg import norm |
|
|
|
cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b)) |
|
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-code', trust_remote_code=True) # trust_remote_code is needed to use the encode method |
|
embeddings = model.encode(['How is the weather today?', 'What is the current weather like today?']) |
|
print(cos_sim(embeddings[0], embeddings[1])) |
|
``` |
|
|
|
If you only want to handle shorter sequence, such as 2k, pass the `max_length` parameter to the `encode` function: |
|
|
|
```python |
|
embeddings = model.encode( |
|
['Very long ... document'], |
|
max_length=2048 |
|
) |
|
``` |
|
|
|
## Fully-managed Embeddings Service |
|
|
|
Alternatively, you can use Jina AI's [Embeddings platform](https://jina.ai/embeddings/) for fully-managed access to Jina Embeddings models. |
|
|
|
## Plans |
|
|
|
The development of new bilingual models is currently underway. We will be targeting mainly the German and Spanish languages. |
|
The upcoming models will be called `jina-embeddings-v2-small-de/es`. |
|
|
|
## Contact |
|
|
|
Join our [Discord community](https://discord.jina.ai) and chat with other community members about ideas. |