Spaces:
Sleeping
Sleeping
File size: 3,538 Bytes
6912ad8 |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from fastembed import TextEmbedding
from fastembed.common.model_description import PoolingType, ModelSource
from huggingface_hub import snapshot_download
import time
# def download_model_from_hf(model_name: str, save_path: str):
# try:
# snapshot_download(
# repo_id=model_name,
# local_dir=save_path,
# allow_patterns=["onnx/*"],
# local_dir_use_symlinks=False,
# )
# except Exception as e:
# print(f"Error downloading model from Hugging Face: {e}")
# raise e
def add_custom_embedding_model(
model_name: str, source_model: str, source_file: str, dim: int, from_hf: bool = True
):
"""Add a custom embedding model to FastEmbed and return an instance of TextEmbedding."""
if from_hf:
try:
TextEmbedding.add_custom_model(
model=model_name,
pooling=PoolingType.MEAN,
normalization=True,
sources=ModelSource(hf=source_model), # can be used with an `url` to load files from a private storage
dim=dim,
model_file=source_file, # can be used to load an already supported model with another optimization or quantization, e.g. onnx/model_O4.onnx
)
print(f"Successfully added model '{model_name}' from Hugging Face.")
return TextEmbedding(model_name=model_name)
except Exception as e:
print(f"Error adding model from Hugging Face: {e}")
raise e
else:
try:
TextEmbedding.add_custom_model(
model=model_name,
pooling=PoolingType.MEAN,
normalization=True,
sources=ModelSource(url=source_model),
dim=dim,
model_file=source_file,
)
print(f"Successfully added model '{model_name}' from local file.")
return TextEmbedding(model_name=model_name)
except Exception as e:
print(f"Error adding model from local file: {e}")
raise e
if __name__ == "__main__":
# Example usage: adding a custom model from Hugging Face
# add_custom_embedding_model(
# model_name="models/Vietnamese_Embedding",
# source_model="AITeamVN/Vietnamese_Embedding",
# source_file="onnx/model.onnx_data",
# dim=1024,
# from_hf=True
# )
# model = TextEmbedding(model_name="AITeamVN/Vietnamese_Embedding")
# embeddings = list(model.embed("text to embed"))
# # Ex
# download_model_from_hf("AITeamVN/Vietnamese_Embedding", "./models/Vietnamese_Embedding")
# from fastembed import TextEmbedding
# from fastembed.common.model_description import PoolingType, ModelSource
TextEmbedding.add_custom_model(
model="Mint1456/Vietnamese_Embedding_OnnX_Quantized",
pooling=PoolingType.MEAN,
normalization=True,
sources=ModelSource(hf="Mint1456/Vietnamese_Embedding_OnnX_Quantized"), # can be used with an `url` to load files from a private storage
dim=1024,
model_file="model.onnx", # can be used to load an already supported model with another optimization or quantization, e.g. onnx/model_O4.onnx
)
model = TextEmbedding(model_name="Mint1456/Vietnamese_Embedding_OnnX_Quantized")
start = time.perf_counter()
embeddings = list(model.embed("define artificial intelligence"))
print(f"len embeding {len(embeddings[0])}, time taken: {time.perf_counter() - start} seconds") |