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

原创

点火三周

发布于 2024-12-06 15:12:28

39200

代码可运行

文章被收录于专栏：Elastic Stack专栏Elastic Stack专栏

运行总次数：0

代码可运行

本文介绍了如何配置和使用Jina Embeddings v2，这是第一个具有8K上下文长度的开源嵌入模型。通过结合Elasticsearch和semantic_text字段类型，我们展示了如何实现后期分块来优化长上下文处理。在详细的步骤中，我们涵盖了创建端点、索引、数据索引、提问和后期分块示例。该方法不仅提高了上下文感知能力，还优化了计算资源的使用，最终提升了语义搜索的精度和效率。

在这篇文章中，我们将配置并使用 jina-embeddings-v2，这是第一个开源的8K上下文长度嵌入模型。我们会从使用 semantic_text 的开箱即用实现开始，然后介绍如何实现后期分块。

长上下文模型

一般情况下，我们使用的嵌入模型的上下文长度是512个标记（tokens），这意味着如果我们尝试创建更长的嵌入，只有前512个标记会被添加到向量字段中。这种短上下文的问题在于，文本块（chunks）只能感知块内的文本，而无法了解整个上下文：

如图所示，在 Chunk 1 中，我们知道在讨论Sarah Johnson，但在 Chunk 2 中，我们失去了直接的引用。因此，随着文档变长，它可能会错过Sarah Johnson首次提到的依赖关系，并且不会将“Sarah Johnson”，“她”和“她的”联系到同一个人。这在有多个人被称为她/她时会变得更加复杂，但我们先来看如何解决这个问题。

传统的长上下文模型只关心前面的词的依赖关系，因为生成文本的任务是根据输入生成下一个词。然而，Jina Embeddings 2模型通过三个关键阶段进行训练：首先，它使用包含1700亿词的英文C4数据集进行掩码词预训练。接下来，它使用Jina AI的新语料库进行成对对比训练，细化嵌入，使相似的文本更加接近，不相似的文本更远。最后，它通过包含相反语法极性的句子的文本三元组和负采样数据集进行微调，以改进处理可能相近但意义相反的句子的能力。

让我们看看这个模型是如何工作的：更长的上下文长度允许我们在同一个块中保留第一次提到Sarah Johnson的引用：

然而，这也有其缺点。上下文越大，意味着你在同一个维度空间中放入了更多的信息。这种压缩可能会稀释上下文，从嵌入中移除潜在的重要信息。另一个缺点是生成更长的嵌入需要更多的计算资源。最后，在一个RAG系统中，文本块的大小决定了你发送给LLM的信息量，这会影响精度、成本和延迟。好消息是你不必使用整个8K标记，你可以根据你的用例找到一个最佳平衡点。

为了结合两者的优点，Jina提出了一种叫做后期分块的方法。后期分块的做法是在生成嵌入后再进行分块，而不是先分块文本，然后为每个独立的块生成嵌入。为此，你需要一个能够创建上下文感知嵌入的模型，然后你可以在保持上下文的同时分块生成的嵌入，即保持块之间的依赖关系和关联。

我们将设置 jina-embeddings-v2 模型在Elasticsearch中使用，并结合 semantic_text ，然后创建一个自定义的后期分块设置。

步骤

创建端点

通过我们的 HuggingFace Open Inference Service 集成，运行 HuggingFace 模型非常简单。你只需打开模型的网页，点击 Inference API 下的 View Code，并从那里获取 API URL。在同一页面，你可以 管理你的令牌 来创建 API Key。

关于创建安全令牌的更多细节，可以访问这里。对于本文的目的，将其设置为 read 令牌即可。

一旦你有了 url 和 api_key，就可以创建推理端点：

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中创建我们自己的后期分块实现。这个过程需要以下步骤：

创建映射

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

加载数据

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

我们不使用ingest pipeline方法，因为我们希望创建特殊的嵌入，而是使用一个Python脚本，其主要作用是获取块标记位置的注释，为整个文档生成嵌入，然后根据我们提供的长度分块嵌入：

通过以下代码，你可以定义文本块大小，通过句子分割并获取块位置。

def chunk_by_sentences(input_text: str, tokenizer: callable):
    """
    将输入文本拆分为句子使用tokenizer
    :param input_text: 要拆分为句子的文本片段
    :param tokenizer: 要使用的tokenizer
    :return: 包含文本块列表及其对应的标记范围的元组
    """
    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

这个第二个函数将接收注释和整个输入的嵌入来生成嵌入块。

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
        ):
            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]

在此过程之后，我们可以索引我们的文档：

# 准备要索引的文档
documents = []
for chunk, new_embedding in zip(chunks, embeddings):
    documents.append(
        {
            "_index": "jina-late-chunking",
            "_source": {
                "content_embedding": new_embedding,
                "content": chunk,
            },
        }
    )
# 使用 helpers.bulk 进行索引
helpers.bulk(client, documents)

你可以在这个 Notebook 中找到完整的示例步骤。

可以自由地在 input_text 变量中尝试不同的值。

运行查询

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

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."
  }
}