Spaces:
Sleeping
Sleeping
File size: 2,053 Bytes
cc74372 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
from transformers import AutoModel, AutoTokenizer
from typing import Any, List, Mapping, Optional,Union
from langchain.callbacks.manager import (
CallbackManagerForLLMRun
)
from langchain_core.embeddings import Embeddings
import torch
class Embedding(Embeddings):
def __init__(self,**kwargs):
self.model=AutoModel.from_pretrained('BAAI/bge-small-zh-v1.5')
self.tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-zh-v1.5')
self.model.eval()
@property
def _llm_type(self) -> str:
return "BAAI/bge-small-zh-v1.5"
@property
def model_name(self) -> str:
return "embedding"
def _call(
self,
prompt: List[str],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
encoded_input = self.tokenizer(prompt, padding=True, truncation=True, return_tensors='pt')
with torch.no_grad():
model_output = self.model(**encoded_input)
# Perform pooling. In this case, cls pooling.
sentence_embeddings = model_output[0][:, 0]
print(sentence_embeddings.shape)
# normalize embeddings
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
return sentence_embeddings.numpy()
@property
def _identifying_params(self) -> Mapping[str, Any]:
"""Get the identifying parameters."""
return {"model_path": self.model_path}
def embed_documents(self, texts) -> List[List[float]]:
# Embed a list of documents
embeddings = []
print("embed_documents:",len(texts),type(texts))
embedding = self._call(texts)
for row in embedding:
embeddings.append(row)
# print("embed_documents: shape",embeddings.shape)
return embeddings
def embed_query(self, text) -> List[float]:
# Embed a single query
embedding = self._call([text])
return embedding[0]
|