From c29cd7ac245a83b317716ad46f07f034188204d4 Mon Sep 17 00:00:00 2001
From: zhangshaohu <shaohu.zhang@fit2cloud.com>
Date: Tue, 30 Jan 2024 19:01:26 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E6=94=B9=E4=B8=8B=E8=BD=BD?=
 =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E8=84=9A=E6=9C=AC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 install_model.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/install_model.py b/install_model.py
index 22394b78d..fb464611a 100644
--- a/install_model.py
+++ b/install_model.py
@@ -8,11 +8,10 @@
 """
 import json
 import os.path
-
+from pycrawlers import huggingface
 from transformers import GPT2TokenizerFast
-import sentence_transformers
-
-prefix_dir = "/opt/maxkb/model"
+hg = huggingface()
+prefix_dir = "./model"
 model_config = [
     {
         'download_params': {
@@ -51,10 +50,10 @@ model_config = [
     },
     {
         'download_params': {
-            'model_name_or_path': 'shibing624/text2vec-base-chinese',
-            'cache_folder': os.path.join(prefix_dir, 'embedding')
+            'urls': ["https://huggingface.co/shibing624/text2vec-base-chinese/tree/main"],
+            'file_save_paths': [os.path.join(prefix_dir, 'embedding',"shibing624_text2vec-base-chinese")]
         },
-        'download_function': sentence_transformers.SentenceTransformer
+        'download_function': hg.get_batch_data
     }
 
 ]