使用 Jina Embeddings v2 在 Elasticsearch 中进行后期分块

作者：来自 Elastic Gustavo Llermaly

在 Elasticsearch 中使用 Jina Embeddings v2 模型并探索长上下文嵌入模型的优缺点。

在本文中，我们将配置和使用 jina-embeddings-v2，这是第一个开源 8K 上下文长度嵌入模型，首先使用 semantic_text 进行 OOTB 实现，然后实现 Late Chunking。

长上下文模型 - long-context models

我们通常看到上下文长度为 512 个 token 的嵌入模型，这意味着如果我们尝试创建更长的嵌入，则只有前 512 个 token 会添加到向量字段中。这些短上下文的问题在于，块不会知道整个上下文，而只会知道块内的文本：

正如你在图片中看到的，在块 1 中我们知道我们在谈论 Sarah Johnson，但在块 2 中我们失去了直接引用。因此，随着文档变长，它可能会错过 Sarah Johnson 首次被提及时的依赖关系，并且不会将 “Sarah Johnson”、“She” 和 “her” 指代同一个人联系起来。当然，如果有不止一个人被称为 her/she，这会变得更加复杂，但现在让我们看看解决这个问题的第一种方法。

旨在生成文本的传统长上下文模型只关心对前面单词的依赖关系，因此输入中的最后一个标记比前面的标记更重要，因为文本生成器的任务是在输入后生成下一个单词。然而，Jina Embeddings 2 模型经过三个关键阶段的训练：首先，它使用 1700 亿个单词的英语 C4 数据集进行掩码单词预训练。接下来，它使用已知相似或不相似的文本对进行成对对比训练，使用 Jina AI 的新语料库来优化嵌入，使相似的文本更接近，而不相似的文本更远。最后，使用文本三元组和负挖掘对其进行微调，结合具有相反语法极性的句子的数据集，以改进对具有相反含义的句子的嵌入可能过于接近的情况的处理。

那么，让我们看看它是如何工作的：更长的上下文长度使我们能够将第一次提到 Sarah Johnson 的引用保留在同一块中：

然而，这也有其缺点。上下文越大，意味着你将在相同维度空间内放置更多信息。这种压缩可能会稀释上下文，从嵌入中删除潜在的重要信息。另一个缺点是生成更长的嵌入需要更多的计算资源。最后，在 RAG 系统中，文本块的大小决定了你向 LLM 发送的信息量，这将影响精度、成本和延迟。好消息是你不必使用整个 8K token，你可以根据你的用例找到一个最佳点。你在文章 “Elasticsearch：检索增强生成背后的重要思想” 可以。

Jina 致力于将两者的优点结合起来，提出了一种称为 “后期分块（Late Chunking）” 的方法。后期分块包括在嵌入之后对文本进行分块，而不是先对文本进行分块，然后为每个独立的块创建嵌入。为此，你需要一个能够创建上下文感知嵌入的模型，然后你可以在保留上下文（即块之间的依赖关系和关系）的同时对生成的嵌入进行分块。

我们将在 Elasticsearch 中设置 jina-embeddings-v2 模型并将其与 semantic_text一起使用，然后为后期分块创建自定义设置。

步骤

创建端点
创建索引
索引数据
提问
后期分块示例

创建端点

借助我们的 HuggingFace 开放推理服务集成（Open Inference Service integration），运行 HuggingFace 模型非常简单。你只需打开模型网页，单击 Inference API 下的 View Code，然后从那里获取 API URL。在同一屏幕中，你可以 Manage your tokens 以创建 API 密钥。

有关创建安全 token 的更多详细信息，你可以访问此处。出于本文的目的，将其设置为 read token 是可以的。

获得 url 和 api_key 后，继续创建推理端点（inference endpoint）：

PUT _inference/text_embedding/jina-embeddings-v2-base-en
{
  "service": "hugging_face",
  "service_settings": {
    "api_key": "hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx", 
    "url": "https://api-inference.huggingface.co/models/jinaai/jina-embeddings-v2-base-en" 
  }
}

如果你收到此错误 “Model jinaai/jina-embeddings-v2-base-en is currently loading”，则表示模型正在预热。请等待几秒钟，然后重试。

创建索引

我们将使用 semantic_text 字段类型。它将负责推断嵌入映射和配置，并为你进行段落分块！如果你想了解更多信息，可以阅读这篇精彩的文章。

PUT jina-embeddings
{
  "mappings": {
    "properties": {
      "super_body": {
        "type": "semantic_text",
        "inference_id": "jina-embeddings-v2-base-en"
      }
    }
  }
}

这种方法将为我们处理向量配置和文档分块，从而为我们提供一个良好的起点。它将创建 250 个单词的块，其中有 100 个单词重叠。对于诸如增加块大小以利用 8K 上下文大小之类的自定义，我们必须经历一个更长的过程，我们将在后期分块部分进行探讨。

索引数据

使用 semantic_text 时，我们会用到。我们只需像往常一样索引数据即可。

PUT jina-embeddings/_bulk
{ "index" : { "_index" : "jina-embeddings", "_id" : "1" } }
{"super_body": "Sarah Johnson is a talented marine biologist working at the Oceanographic Institute. Her groundbreaking research on coral reef ecosystems has garnered international attention and numerous accolades."}
{ "index" : { "_index" : "jina-embeddings", "_id" : "2" } }
{"super_body": "She spends months at a time diving in remote locations, meticulously documenting the intricate relationships between various marine species. "}
{ "index" : { "_index" : "jina-embeddings", "_id" : "3" } }
{"super_body": "Her dedication to preserving these delicate underwater environments has inspired a new generation of conservationists."}

提出问题

现在我们可以使用语义搜索查询来向我们的数据提出问题：

GET jina-embeddings/_search 
{
  "query": {
    "semantic": {
      "field": "super_body",
      "query": "who inspired taking care of the sea?"
    }
  }
}

第一个结果将如下所示：

{
    "_index": "jina-embeddings",
    "_id": "1",
    "_score": 0.64889884,
    "_source": {
        "super_body": {
            "text": "Sarah Johnson is a talented marine biologist working at the Oceanographic Institute. Her groundbreaking research on coral reef ecosystems has garnered international attention and numerous accolades.",
            "inference": {
                "inference_id": "jina-embeddings-v2-base-en",
                "model_settings": {
                    "task_type": "text_embedding",
                    "dimensions": 768,
                    "similarity": "cosine",
                    "element_type": "float"
                },
                "chunks": [
                    {
                        "text": "Sarah Johnson is a talented marine biologist working at the Oceanographic Institute. Her groundbreaking research on coral reef ecosystems has garnered international attention and numerous accolades.",
                        "embeddings": [
                            -0.0064849486,
                            -0.014192865,
                            0.028806737,
                            0.0026694024,
                            ... // 768 dims
                        ]
                    }
                ]
            }
        }
    }
}

后期分块示例

现在我们已经配置了嵌入模型，我们可以在 Elasticsearch 中创建自己的后期分块实现。该过程需要以下步骤：

1. 创建映射

PUT jina-late-chunking
{
  "mappings": {
    "properties": {
      "content_embedding": { 
        "type": "dense_vector", 
        "dims": 768, 
        "element_type": "float",
        "similarity": "cosine"
      },
      "content": { 
        "type": "text" 
      }
    }
  }
}

2. 加载数据

你可以在支持 Notebook 中找到完整的实现。

我们在这里不使用摄取管道方法，因为我们想要创建特殊的嵌入，而是使用一个 Python 脚本，其关键作用是获取块标记位置的注释，为整个文档生成嵌入，然后根据我们提供的长度对嵌入进行分块：

使用此代码，你可以通过按句子拆分并获取块位置来定义文本块大小。

def chunk_by_sentences(input_text: str, tokenizer: callable):
    """
    Split the input text into sentences using the tokenizer
    :param input_text: The text snippet to split into sentences
    :param tokenizer: The tokenizer to use
    :return: A tuple containing the list of text chunks and their corresponding token spans
    """
    inputs = tokenizer(input_text, return_tensors='pt', return_offsets_mapping=True)
    punctuation_mark_id = tokenizer.convert_tokens_to_ids('.')
    sep_id = tokenizer.convert_tokens_to_ids('[SEP]')
    token_offsets = inputs['offset_mapping'][0]
    token_ids = inputs['input_ids'][0]
    chunk_positions = [
        (i, int(start + 1))
        for i, (token_id, (start, end)) in enumerate(zip(token_ids, token_offsets))
        if token_id == punctuation_mark_id
        and (
            token_offsets[i + 1][0] - token_offsets[i][1] > 0
            or token_ids[i + 1] == sep_id
        )
    ]
    chunks = [
        input_text[x[1] : y[1]]
        for x, y in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)
    ]
    span_annotations = [
        (x[0], y[0]) for (x, y) in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)
    ]
    return chunks, span_annotations

第二个函数将接收整个输入的注释（annotations）和嵌入以生成嵌入块。

def late_chunking(
    model_output: 'BatchEncoding', span_annotation: list, max_length=None
):
    token_embeddings = model_output[0]
    outputs = []
    for embeddings, annotations in zip(token_embeddings, span_annotation):
        if (
            max_length is not None
        ):  # remove annotations which go bejond the max-length of the model
            annotations = [
                (start, min(end, max_length - 1))
                for (start, end) in annotations
                if start < (max_length - 1)
            ]
        pooled_embeddings = [
            embeddings[start:end].sum(dim=0) / (end - start)
            for start, end in annotations
            if (end - start) >= 1
        ]
        pooled_embeddings = [
            embedding.detach().cpu().numpy() for embedding in pooled_embeddings
        ]
        outputs.append(pooled_embeddings)

    return outputs

这是将所有内容组合在一起的部分；对整个文本输入进行标记，然后将其传递给 late_chunking 函数以对池化嵌入进行分块。

inputs = tokenizer(input_text, return_tensors='pt')
model_output = model(**inputs)
embeddings = late_chunking(model_output, [span_annotations])[0]

经过这个过程，我们可以索引我们的文档：

# Prepare the documents to be indexed
documents = []
for chunk, new_embedding in zip(chunks, embeddings):
    documents.append(
        {
            "_index": "jina-late-chunking",
            "_source": {
                "content_embedding": new_embedding,
                "content": chunk,
            },
        }
    )
# Use helpers.bulk to index
helpers.bulk(client, documents)

你可以在此处找到包含完整示例的笔记本。

请随意尝试 input_text 变量中的不同值。

3. 运行查询

你现在可以针对新数据索引运行语义搜索：

GET jina-late-chunking/_search
{
  "knn": {
    "field": "content_embedding",
    "query_vector_builder": {
      "text_embedding": {
        "model_id": "jina-embeddings-v2-base-en",
        "model_text": "berlin"
      }
    },
    "k": 10,
    "num_candidates": 100
  }
}

结果将如下所示：

{
  "_index": "jina-late-chunking",
  "_id": "gGDN1JEBF7lnCNFTVZBg",
  "_score": 0.4930191,
  "_source": {
    "content_embedding": [
      -0.9107036590576172,
      -0.57366544008255,
      1.0492067337036133,
      0.25255489349365234,
      -0.1283145546913147... 
    ],
    "content": "Berlin is the capital and largest city of Germany, both by area and by population."
  }
}

结论

虽然仍处于试验阶段，但后期分块可能有很多好处，尤其是在 RAG 中，因为它允许你在对文本进行分块时保留关键上下文信息。此外，Jina 嵌入模型有助于存储较短的向量，从而占用更少的内存和存储空间，并加快搜索检索速度。因此，这两个功能与 Elasticsearch 结合使用，在使用向量搜索时提高了管理和检索信息的效率和有效性。

Elasticsearch 与业界领先的 Gen AI 工具和提供商进行了原生集成。查看我们的网络研讨会，了解如何超越 RAG 基础知识，或构建可用于生产的应用程序 Elastic Vector Database。

要为你的用例构建最佳搜索解决方案，请立即开始免费云试用或在你的本地机器上试用 Elastic。

原文：Late chunking in Elasticsearch with Jina Embeddings v2 - Search Labs