wehub-resource-sync
31 lines
1.3 KiB
Python
31 lines
1.3 KiB
Python
# !/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
# @author: CS_木成河
|
|
# @time: 2024/10/31 14:33
|
|
# @blog: https://blog.csdn.net/weixin_47936614
|
|
|
|
from langchain_core.embeddings import Embeddings
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
|
|
class RagTextEmbeddings(Embeddings):
|
|
def __init__(self, embed_model_path: str, **kwargs):
|
|
self.batch_size = kwargs['batch_size']
|
|
self.device = kwargs['device']
|
|
self.embed_model = SentenceTransformer(embed_model_path, trust_remote_code=True, device=self.device)
|
|
|
|
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
|
docs_embeddings = self.embed_model.encode(texts,
|
|
task="retrieval.passage",
|
|
batch_size=self.batch_size,
|
|
device=self.device,
|
|
show_progress_bar=True)
|
|
return docs_embeddings.tolist()
|
|
|
|
def embed_query(self, text: str) -> list[float]:
|
|
query_embeddings = self.embed_model.encode([text],
|
|
task="retrieval.query",
|
|
device=self.device)
|
|
return query_embeddings.tolist()[0]
|
|
|