From cc1e8449900fa25ace46f834ad4bee4adc3d427c Mon Sep 17 00:00:00 2001 From: Kollin Date: Mon, 22 Jul 2024 10:04:20 +0800 Subject: [PATCH] support safetensors BGE model support safetensors --- .../src/pybackend_libs/dataelem/model/embedding/bge.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/pybackend_libs/src/pybackend_libs/dataelem/model/embedding/bge.py b/python/pybackend_libs/src/pybackend_libs/dataelem/model/embedding/bge.py index fb26deb..c678253 100644 --- a/python/pybackend_libs/src/pybackend_libs/dataelem/model/embedding/bge.py +++ b/python/pybackend_libs/src/pybackend_libs/dataelem/model/embedding/bge.py @@ -13,6 +13,7 @@ def __init__(self, **kwargs): precision = kwargs.get('precision', 'fp16') gpu_memory = kwargs.get('gpu_memory') devices = kwargs.get('devices').split(',') + use_safetensors = True if kwargs.get("use_safetensors") == "1" else False self.devices = devices self.default_device = f'cuda:{devices[0]}' self.batch_size = int(kwargs.get('batch_size', '32')) @@ -25,6 +26,7 @@ def __init__(self, **kwargs): precision, devices, gpu_memory, + use_safetensors=use_safetensors, ) def predict(self, kwargs):