feat: 模型下载脚本

apps/application/serializers/chat_message_serializers.py

@@ -140,7 +140,9 @@ class ChatMessageSerializer(serializers.Serializer):
                     ChatMessage(_id, message, title, content, embedding_id, dataset_id, document_id,
                                 paragraph_id,
                                 source_type,
-                                source_id, c, 0,
+                                source_id,
+                                c if c is not None else '抱歉，根据已知信息无法回答这个问题，请重新描述您的问题或提供更多信息～',
+                                0,
                                 0))
                 # 重新设置缓存
                 chat_cache.set(chat_id,

apps/smartdoc/settings/base.py

@@ -157,7 +157,7 @@ TIME_ZONE = CONFIG.get_time_zone()
 
 USE_I18N = True
 
-USE_TZ = True
+USE_TZ = False
 
 # Static files (CSS, JavaScript, Images)
 # https://docs.djangoproject.com/en/4.2/howto/static-files/

install_model.py

@@ -0,0 +1,71 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： install_model.py
+    @date：2023/12/18 14:02
+    @desc:
+"""
+import json
+import os.path
+
+from transformers import GPT2TokenizerFast
+import sentence_transformers
+
+prefix_dir = os.path.join(os.path.dirname(os.path.abspath(os.getcwd())), 'model')
+
+model_config = [
+    {
+        'download_params': {
+            'cache_dir': os.path.join(prefix_dir, 'base'),
+            'pretrained_model_name_or_path': 'gpt2'
+        },
+        'download_function': GPT2TokenizerFast.from_pretrained
+    },
+    {
+        'download_params': {
+            'cache_dir': os.path.join(prefix_dir, 'base'),
+            'pretrained_model_name_or_path': 'gpt2-medium'
+        },
+        'download_function': GPT2TokenizerFast.from_pretrained
+    },
+    {
+        'download_params': {
+            'cache_dir': os.path.join(prefix_dir, 'base'),
+            'pretrained_model_name_or_path': 'gpt2-large'
+        },
+        'download_function': GPT2TokenizerFast.from_pretrained
+    },
+    {
+        'download_params': {
+            'cache_dir': os.path.join(prefix_dir, 'base'),
+            'pretrained_model_name_or_path': 'gpt2-xl'
+        },
+        'download_function': GPT2TokenizerFast.from_pretrained
+    },
+    {
+        'download_params': {
+            'cache_dir': os.path.join(prefix_dir, 'base'),
+            'pretrained_model_name_or_path': 'distilgpt2'
+        },
+        'download_function': GPT2TokenizerFast.from_pretrained
+    },
+    {
+        'download_params': {
+            'model_name_or_path': 'shibing624/text2vec-base-chinese',
+            'cache_folder': os.path.join(prefix_dir, 'embedding')
+        },
+        'download_function': sentence_transformers.SentenceTransformer
+    }
+
+]
+
+
+def install():
+    for model in model_config:
+        print(json.dumps(model.get('download_params')))
+        model.get('download_function')(**model.get('download_params'))
+
+
+if __name__ == '__main__':
+    install()

main.py

@@ -48,6 +48,7 @@ def start_services():
 
 
 if __name__ == '__main__':
+    os.environ['TRANSFORMERS_CACHE'] = '/opt/maxkb/model'
     parser = argparse.ArgumentParser(
         description="""
            qabot service control tools;