diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 000000000..a2998b85e --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,38 @@ +--- +name: Bug report +about: Create a report to help us improve +title: "[BUG]: " +labels: '' +assignees: '' + +--- + +**Describe the bug** +A clear and concise description of what the bug is. + +**To Reproduce** +Steps to reproduce the behavior: +1. Go to '...' +2. Click on '....' +3. Scroll down to '....' +4. See error + +**Expected behavior** +A clear and concise description of what you expected to happen. + +**Screenshots** +If applicable, add screenshots to help explain your problem. + +**Desktop (please complete the following information):** + - OS: [e.g. iOS] + - Browser [e.g. chrome, safari] + - Version [e.g. 22] + +**Smartphone (please complete the following information):** + - Device: [e.g. iPhone6] + - OS: [e.g. iOS8.1] + - Browser [e.g. stock browser, safari] + - Version [e.g. 22] + +**Additional context** +Add any other context about the problem here. diff --git a/.github/ISSUE_TEMPLATE/documentation-related.md b/.github/ISSUE_TEMPLATE/documentation-related.md new file mode 100644 index 000000000..5506434f7 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/documentation-related.md @@ -0,0 +1,10 @@ +--- +name: Documentation Related +about: Describe this issue template's purpose here. +title: "[Doc]: " +labels: '' +assignees: '' + +--- + + diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 000000000..0ccd363af --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,20 @@ +--- +name: Feature request +about: Suggest an idea for this project +title: "[Feature]:" +labels: '' +assignees: '' + +--- + +**Is your feature request related to a problem? Please describe.** +A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] + +**Describe the solution you'd like** +A clear and concise description of what you want to happen. + +**Describe alternatives you've considered** +A clear and concise description of any alternative solutions or features you've considered. + +**Additional context** +Add any other context or screenshots about the feature request here. diff --git a/.gitignore b/.gitignore index 5043f7db0..22bb204db 100644 --- a/.gitignore +++ b/.gitignore @@ -23,6 +23,7 @@ lib/ lib64/ parts/ sdist/ +models var/ wheels/ models/ diff --git a/README.md b/README.md index cd7aaff9b..18dd30f80 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,10 @@ Currently, we have released multiple key features, which are listed below to dem - Unified vector storage/indexing of knowledge base - Support for unstructured data such as PDF, Markdown, CSV, and WebURL +- Milti LLMs Support + - Supports multiple large language models, currently supporting Vicuna (7b, 13b), ChatGLM-6b (int4, int8) + - TODO: codegen2, codet5p + ## Demo @@ -177,6 +181,10 @@ Notice: the webserver need to connect llmserver, so you need change the .env f We provide a user interface for Gradio, which allows you to use DB-GPT through our user interface. Additionally, we have prepared several reference articles (written in Chinese) that introduce the code and principles related to our project. - [LLM Practical In Action Series (1) — Combined Langchain-Vicuna Application Practical](https://medium.com/@cfqcsunny/llm-practical-in-action-series-1-combined-langchain-vicuna-application-practical-701cd0413c9f) +### Multi LLMs Usage + +To use multiple models, modify the LLM_MODEL parameter in the .env configuration file to switch between the models. + ####Create your own knowledge repository: 1.Place personal knowledge files or folders in the pilot/datasets directory. @@ -215,7 +223,7 @@ The achievements of this project are thanks to the technical community, especial | :---: | :---: | :---: | :---: |:---: | -This project follows the git-contributor [spec](https://github.com/xudafeng/git-contributor), auto updated at `Sun May 14 2023 23:02:43 GMT+0800`. +This project follows the git-contributor [spec](https://github.com/xudafeng/git-contributor), auto updated at `Fri May 19 2023 00:24:18 GMT+0800`. diff --git a/README.zh.md b/README.zh.md index 2c6ecca43..b84671da2 100644 --- a/README.zh.md +++ b/README.zh.md @@ -26,6 +26,10 @@ DB-GPT 是一个开源的以数据库为基础的GPT实验项目，使用本地 - 知识库统一向量存储/索引 - 非结构化数据支持包括PDF、MarkDown、CSV、WebURL +- 多模型支持 + - 支持多种大语言模型, 当前已支持Vicuna(7b,13b), ChatGLM-6b(int4, int8) + - TODO: codet5p, codegen2 + ## 效果演示 示例通过 RTX 4090 GPU 演示，[YouTube 地址](https://www.youtube.com/watch?v=1PWI6F89LPo) @@ -180,6 +184,10 @@ $ python webserver.py 2. [大模型实战系列(2) —— DB-GPT 阿里云部署指南](https://zhuanlan.zhihu.com/p/629467580) 3. [大模型实战系列(3) —— DB-GPT插件模型原理与使用](https://zhuanlan.zhihu.com/p/629623125) + +### 多模型使用 +在.env 配置文件当中, 修改LLM_MODEL参数来切换使用的模型。 + ####打造属于你的知识库： 1、将个人知识文件或者文件夹放入pilot/datasets目录中 @@ -212,12 +220,14 @@ python tools/knowledge_init.py +## 贡献者 ## Contributors |[
csunny](https://github.com/csunny)
|[
xudafeng](https://github.com/xudafeng)
|[
明天](https://github.com/yhjun1026)
| [
Aries-ckt](https://github.com/Aries-ckt)
|[
thebigbone](https://github.com/thebigbone)
| | :---: | :---: | :---: | :---: |:---: | +[git-contributor 说明](https://github.com/xudafeng/git-contributor)，自动生成时间：`Fri May 19 2023 00:24:18 GMT+0800`。 This project follows the git-contributor [spec](https://github.com/xudafeng/git-contributor), auto updated at `Sun May 14 2023 23:02:43 GMT+0800`. diff --git a/examples/embdserver.py b/examples/embdserver.py index 79140ba66..32eca1291 100644 --- a/examples/embdserver.py +++ b/examples/embdserver.py @@ -5,14 +5,21 @@ import requests import json import time import uuid +import os +import sys from urllib.parse import urljoin import gradio as gr + +ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.append(ROOT_PATH) + + from pilot.configs.config import Config from pilot.conversation import conv_qa_prompt_template, conv_templates from langchain.prompts import PromptTemplate -vicuna_stream_path = "generate_stream" +llmstream_stream_path = "generate_stream" CFG = Config() @@ -21,38 +28,45 @@ def generate(query): template_name = "conv_one_shot" state = conv_templates[template_name].copy() - pt = PromptTemplate( - template=conv_qa_prompt_template, - input_variables=["context", "question"] - ) + # pt = PromptTemplate( + # template=conv_qa_prompt_template, + # input_variables=["context", "question"] + # ) - result = pt.format(context="This page covers how to use the Chroma ecosystem within LangChain. It is broken into two parts: installation and setup, and then references to specific Chroma wrappers.", - question=query) + # result = pt.format(context="This page covers how to use the Chroma ecosystem within LangChain. It is broken into two parts: installation and setup, and then references to specific Chroma wrappers.", + # question=query) - print(result) + # print(result) - state.append_message(state.roles[0], result) + state.append_message(state.roles[0], query) state.append_message(state.roles[1], None) prompt = state.get_prompt() params = { - "model": "vicuna-13b", + "model": "chatglm-6b", "prompt": prompt, - "temperature": 0.7, + "temperature": 1.0, "max_new_tokens": 1024, "stop": "###" } response = requests.post( - url=urljoin(CFG.MODEL_SERVER, vicuna_stream_path), data=json.dumps(params) + url=urljoin(CFG.MODEL_SERVER, llmstream_stream_path), data=json.dumps(params) ) skip_echo_len = len(params["prompt"]) + 1 - params["prompt"].count("") * 3 + for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"): + if chunk: data = json.loads(chunk.decode()) if data["error_code"] == 0: - output = data["text"][skip_echo_len:].strip() + + if "vicuna" in CFG.LLM_MODEL: + output = data["text"][skip_echo_len:].strip() + else: + output = data["text"].strip() + state.messages[-1][-1] = output + "▌" yield(output) diff --git a/pilot/configs/config.py b/pilot/configs/config.py index 9023bc061..b914390f7 100644 --- a/pilot/configs/config.py +++ b/pilot/configs/config.py @@ -105,7 +105,8 @@ class Config(metaclass=Singleton): self.LLM_MODEL = os.getenv("LLM_MODEL", "vicuna-13b") self.LIMIT_MODEL_CONCURRENCY = int(os.getenv("LIMIT_MODEL_CONCURRENCY", 5)) self.MAX_POSITION_EMBEDDINGS = int(os.getenv("MAX_POSITION_EMBEDDINGS", 4096)) - self.MODEL_SERVER = os.getenv("MODEL_SERVER", "http://121.41.167.183:8000") + self.MODEL_PORT = os.getenv("MODEL_PORT", 8000) + self.MODEL_SERVER = os.getenv("MODEL_SERVER", "http://127.0.0.1" + ":" + str(self.MODEL_PORT)) self.ISLOAD_8BIT = os.getenv("ISLOAD_8BIT", "True") == "True" def set_debug_mode(self, value: bool) -> None: diff --git a/pilot/configs/model_config.py b/pilot/configs/model_config.py index c63187d03..7c4928304 100644 --- a/pilot/configs/model_config.py +++ b/pilot/configs/model_config.py @@ -16,11 +16,17 @@ DATA_DIR = os.path.join(PILOT_PATH, "data") nltk.data.path = [os.path.join(PILOT_PATH, "nltk_data")] + nltk.data.path -DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" LLM_MODEL_CONFIG = { "flan-t5-base": os.path.join(MODEL_PATH, "flan-t5-base"), "vicuna-13b": os.path.join(MODEL_PATH, "vicuna-13b"), + "vicuna-7b": os.path.join(MODEL_PATH, "vicuna-7b"), "text2vec": os.path.join(MODEL_PATH, "text2vec-large-chinese"), + "sentence-transforms": os.path.join(MODEL_PATH, "all-MiniLM-L6-v2"), + "codegen2-1b": os.path.join(MODEL_PATH, "codegen2-1B"), + "codet5p-2b": os.path.join(MODEL_PATH, "codet5p-2b"), + "chatglm-6b-int4": os.path.join(MODEL_PATH, "chatglm-6b-int4"), + "chatglm-6b": os.path.join(MODEL_PATH, "chatglm-6b"), "text2vec-base": os.path.join(MODEL_PATH, "text2vec-base-chinese"), "sentence-transforms": os.path.join(MODEL_PATH, "all-MiniLM-L6-v2") } @@ -29,7 +35,7 @@ LLM_MODEL_CONFIG = { VECTOR_SEARCH_TOP_K = 20 LLM_MODEL = "vicuna-13b" LIMIT_MODEL_CONCURRENCY = 5 -MAX_POSITION_EMBEDDINGS = 4096 +MAX_POSITION_EMBEDDINGS = 4096 # VICUNA_MODEL_SERVER = "http://121.41.227.141:8000" VICUNA_MODEL_SERVER = "http://120.79.27.110:8000" @@ -38,15 +44,9 @@ ISLOAD_8BIT = True ISDEBUG = False -DB_SETTINGS = { - "user": "root", - "password": "aa123456", - "host": "127.0.0.1", - "port": 3306 -} - +VECTOR_SEARCH_TOP_K = 10 VS_ROOT_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "vs_store") KNOWLEDGE_UPLOAD_ROOT_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data") KNOWLEDGE_CHUNK_SPLIT_SIZE = 100 -VECTOR_STORE_TYPE = "Chroma" +VECTOR_STORE_TYPE = "milvus" VECTOR_STORE_CONFIG = {"url": "127.0.0.1", "port": "19530"} diff --git a/pilot/conversation.py b/pilot/conversation.py index 7f526fb89..0470bc720 100644 --- a/pilot/conversation.py +++ b/pilot/conversation.py @@ -15,6 +15,9 @@ DB_SETTINGS = { "port": CFG.LOCAL_DB_PORT } +ROLE_USER = "USER" +ROLE_ASSISTANT = "Assistant" + class SeparatorStyle(Enum): SINGLE = auto() TWO = auto() diff --git a/pilot/model/adapter.py b/pilot/model/adapter.py new file mode 100644 index 000000000..be8980726 --- /dev/null +++ b/pilot/model/adapter.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +from typing import List +from functools import cache + +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + AutoModel +) + +from pilot.configs.model_config import DEVICE + +class BaseLLMAdaper: + """The Base class for multi model, in our project. + We will support those model, which performance resemble ChatGPT """ + + def match(self, model_path: str): + return True + + def loader(self, model_path: str, from_pretrained_kwargs: dict): + tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) + model = AutoModelForCausalLM.from_pretrained( + model_path, low_cpu_mem_usage=True, **from_pretrained_kwargs + ) + return model, tokenizer + + +llm_model_adapters: List[BaseLLMAdaper] = [] + +# Register llm models to adapters, by this we can use multi models. +def register_llm_model_adapters(cls): + """Register a llm model adapter.""" + llm_model_adapters.append(cls()) + + +@cache +def get_llm_model_adapter(model_path: str) -> BaseLLMAdaper: + for adapter in llm_model_adapters: + if adapter.match(model_path): + return adapter + + raise ValueError(f"Invalid model adapter for {model_path}") + + +# TODO support cpu? for practise we support gpt4all or chatglm-6b-int4? + +class VicunaLLMAdapater(BaseLLMAdaper): + """Vicuna Adapter """ + def match(self, model_path: str): + return "vicuna" in model_path + + def loader(self, model_path: str, from_pretrained_kwagrs: dict): + tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) + model = AutoModelForCausalLM.from_pretrained( + model_path, + low_cpu_mem_usage=True, + **from_pretrained_kwagrs + ) + return model, tokenizer + +class ChatGLMAdapater(BaseLLMAdaper): + """LLM Adatpter for THUDM/chatglm-6b""" + def match(self, model_path: str): + return "chatglm" in model_path + + def loader(self, model_path: str, from_pretrained_kwargs: dict): + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + + if DEVICE != "cuda": + model = AutoModel.from_pretrained( + model_path, trust_remote_code=True, **from_pretrained_kwargs + ).float() + return model, tokenizer + else: + model = AutoModel.from_pretrained( + model_path, trust_remote_code=True, **from_pretrained_kwargs + ).half().cuda() + return model, tokenizer + +class CodeGenAdapter(BaseLLMAdaper): + pass + +class StarCoderAdapter(BaseLLMAdaper): + pass + +class T5CodeAdapter(BaseLLMAdaper): + pass + +class KoalaLLMAdapter(BaseLLMAdaper): + """Koala LLM Adapter which Based LLaMA """ + def match(self, model_path: str): + return "koala" in model_path + + +class RWKV4LLMAdapter(BaseLLMAdaper): + """LLM Adapter for RwKv4 """ + def match(self, model_path: str): + return "RWKV-4" in model_path + + def loader(self, model_path: str, from_pretrained_kwargs: dict): + # TODO + pass + +class GPT4AllAdapter(BaseLLMAdaper): + """A light version for someone who want practise LLM use laptop.""" + def match(self, model_path: str): + return "gpt4all" in model_path + + +register_llm_model_adapters(VicunaLLMAdapater) +register_llm_model_adapters(ChatGLMAdapater) +# TODO Default support vicuna, other model need to tests and Evaluate + +register_llm_model_adapters(BaseLLMAdaper) \ No newline at end of file diff --git a/pilot/model/chat.py b/pilot/model/chat.py deleted file mode 100644 index 97206f2d5..000000000 --- a/pilot/model/chat.py +++ /dev/null @@ -1,3 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding:utf-8 -*- - diff --git a/pilot/model/chatglm_llm.py b/pilot/model/chatglm_llm.py new file mode 100644 index 000000000..0f8b74efa --- /dev/null +++ b/pilot/model/chatglm_llm.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- + +import torch + +from pilot.conversation import ROLE_USER, ROLE_ASSISTANT + +@torch.inference_mode() +def chatglm_generate_stream(model, tokenizer, params, device, context_len=2048, stream_interval=2): + + """Generate text using chatglm model's chat api """ + prompt = params["prompt"] + temperature = float(params.get("temperature", 1.0)) + top_p = float(params.get("top_p", 1.0)) + stop = params.get("stop", "###") + echo = params.get("echo", False) + + generate_kwargs = { + "do_sample": True if temperature > 1e-5 else False, + "top_p": top_p, + "repetition_penalty": 1.0, + "logits_processor": None + } + + if temperature > 1e-5: + generate_kwargs["temperature"] = temperature + + # TODO, Fix this + hist = [] + + messages = prompt.split(stop) + + # Add history chat to hist for model. + for i in range(1, len(messages) - 2, 2): + hist.append((messages[i].split(ROLE_USER + ":")[1], messages[i+1].split(ROLE_ASSISTANT + ":")[1])) + + query = messages[-2].split(ROLE_USER + ":")[1] + print("Query Message: ", query) + output = "" + i = 0 + for i, (response, new_hist) in enumerate(model.stream_chat(tokenizer, query, hist, **generate_kwargs)): + if echo: + output = query + " " + response + else: + output = response + + yield output + + yield output \ No newline at end of file diff --git a/pilot/model/llm/monkey_patch.py b/pilot/model/llm/monkey_patch.py new file mode 100644 index 000000000..a50481281 --- /dev/null +++ b/pilot/model/llm/monkey_patch.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- + +import math +from typing import Optional, Tuple + +import torch +from torch import nn +import transformers + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2].clone() + x2 = x[..., x.shape[-1] // 2 :].clone() + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids): + gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1] + gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3]) + cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) + sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: bool = False, + use_cache: bool = False, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + query_states = ( + self.q_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + key_states = ( + self.k_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + value_states = ( + self.v_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value[0].shape[-2] + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin, position_ids + ) + # [bsz, nh, t, hd] + + if past_key_value is not None: + # reuse k, v, self_attention + key_states = torch.cat([past_key_value[0], key_states], dim=2) + value_states = torch.cat([past_key_value[1], value_states], dim=2) + + past_key_value = (key_states, value_states) if use_cache else None + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt( + self.head_dim + ) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + attn_weights = torch.max( + attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min) + ) + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( + query_states.dtype + ) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +def replace_llama_attn_with_non_inplace_operations(): + """Avoid bugs in mps backend by not using in-place operations.""" + transformers.models.llama.modeling_llama.LlamaAttention.forward = forward + +import transformers + + + +def replace_llama_attn_with_non_inplace_operations(): + """Avoid bugs in mps backend by not using in-place operations.""" + transformers.models.llama.modeling_llama.LlamaAttention.forward = forward diff --git a/pilot/model/loader.py b/pilot/model/loader.py index 66d9c733e..531080314 100644 --- a/pilot/model/loader.py +++ b/pilot/model/loader.py @@ -2,11 +2,39 @@ # -*- coding: utf-8 -*- import torch +import sys import warnings from pilot.singleton import Singleton - +from typing import Optional from pilot.model.compression import compress_module from pilot.model.adapter import get_llm_model_adapter +from pilot.utils import get_gpu_memory +from pilot.model.llm.monkey_patch import replace_llama_attn_with_non_inplace_operations + +def raise_warning_for_incompatible_cpu_offloading_configuration( + device: str, load_8bit: bool, cpu_offloading: bool +): + if cpu_offloading: + if not load_8bit: + warnings.warn( + "The cpu-offloading feature can only be used while also using 8-bit-quantization.\n" + "Use '--load-8bit' to enable 8-bit-quantization\n" + "Continuing without cpu-offloading enabled\n" + ) + return False + if not "linux" in sys.platform: + warnings.warn( + "CPU-offloading is only supported on linux-systems due to the limited compatability with the bitsandbytes-package\n" + "Continuing without cpu-offloading enabled\n" + ) + return False + if device != "cuda": + warnings.warn( + "CPU-offloading is only enabled when using CUDA-devices\n" + "Continuing without cpu-offloading enabled\n" + ) + return False + return cpu_offloading class ModelLoader(metaclass=Singleton): @@ -30,26 +58,37 @@ class ModelLoader(metaclass=Singleton): } # TODO multi gpu support - def loader(self, num_gpus, load_8bit=False, debug=False): + def loader(self, num_gpus, load_8bit=False, debug=False, cpu_offloading=False, max_gpu_memory: Optional[str]=None): + if self.device == "cpu": - kwargs = {} + kwargs = {"torch_dtype": torch.float32} elif self.device == "cuda": kwargs = {"torch_dtype": torch.float16} - if num_gpus == "auto": + num_gpus = int(num_gpus) + + if num_gpus != 1: kwargs["device_map"] = "auto" + if max_gpu_memory is None: + kwargs["device_map"] = "sequential" + + available_gpu_memory = get_gpu_memory(num_gpus) + kwargs["max_memory"] = { + i: str(int(available_gpu_memory[i] * 0.85)) + "GiB" + for i in range(num_gpus) + } + else: - num_gpus = int(num_gpus) - if num_gpus != 1: - kwargs.update({ - "device_map": "auto", - "max_memory": {i: "13GiB" for i in range(num_gpus)}, - }) + kwargs["max_memory"] = {i: max_gpu_memory for i in range(num_gpus)} + + elif self.device == "mps": + kwargs = kwargs = {"torch_dtype": torch.float16} + replace_llama_attn_with_non_inplace_operations() else: - # Todo Support mps for practise raise ValueError(f"Invalid device: {self.device}") - + # TODO when cpu loading, need use quantization config + llm_adapter = get_llm_model_adapter(self.model_path) model, tokenizer = llm_adapter.loader(self.model_path, kwargs) @@ -61,7 +100,7 @@ class ModelLoader(metaclass=Singleton): else: compress_module(model, self.device) - if (self.device == "cuda" and num_gpus == 1): + if (self.device == "cuda" and num_gpus == 1 and not cpu_offloading) or self.device == "mps": model.to(self.device) if debug: diff --git a/pilot/server/chat_adapter.py b/pilot/server/chat_adapter.py new file mode 100644 index 000000000..805cacb3d --- /dev/null +++ b/pilot/server/chat_adapter.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +from typing import List +from functools import cache +from pilot.model.inference import generate_stream + +class BaseChatAdpter: + """The Base class for chat with llm models. it will match the model, + and fetch output from model""" + + def match(self, model_path: str): + return True + + def get_generate_stream_func(self): + """Return the generate stream handler func""" + pass + + +llm_model_chat_adapters: List[BaseChatAdpter] = [] + + +def register_llm_model_chat_adapter(cls): + """Register a chat adapter""" + llm_model_chat_adapters.append(cls()) + + +@cache +def get_llm_chat_adapter(model_path: str) -> BaseChatAdpter: + """Get a chat generate func for a model""" + for adapter in llm_model_chat_adapters: + if adapter.match(model_path): + return adapter + + raise ValueError(f"Invalid model for chat adapter {model_path}") + + +class VicunaChatAdapter(BaseChatAdpter): + + """ Model chat Adapter for vicuna""" + def match(self, model_path: str): + return "vicuna" in model_path + + def get_generate_stream_func(self): + return generate_stream + + +class ChatGLMChatAdapter(BaseChatAdpter): + """ Model chat Adapter for ChatGLM""" + def match(self, model_path: str): + return "chatglm" in model_path + + def get_generate_stream_func(self): + from pilot.model.chatglm_llm import chatglm_generate_stream + return chatglm_generate_stream + + +class CodeT5ChatAdapter(BaseChatAdpter): + + """ Model chat adapter for CodeT5 """ + def match(self, model_path: str): + return "codet5" in model_path + + def get_generate_stream_func(self): + # TODO + pass + +class CodeGenChatAdapter(BaseChatAdpter): + + """ Model chat adapter for CodeGen """ + def match(self, model_path: str): + return "codegen" in model_path + + def get_generate_stream_func(self): + # TODO + pass + + +register_llm_model_chat_adapter(VicunaChatAdapter) +register_llm_model_chat_adapter(ChatGLMChatAdapter) + +register_llm_model_chat_adapter(BaseChatAdpter) \ No newline at end of file diff --git a/pilot/server/llmserver.py b/pilot/server/llmserver.py index e341cc457..bc227d518 100644 --- a/pilot/server/llmserver.py +++ b/pilot/server/llmserver.py @@ -23,20 +23,65 @@ from pilot.model.inference import generate_output, get_embeddings from pilot.model.loader import ModelLoader from pilot.configs.model_config import * from pilot.configs.config import Config +from pilot.server.chat_adapter import get_llm_chat_adapter CFG = Config() -model_path = LLM_MODEL_CONFIG[CFG.LLM_MODEL] - -ml = ModelLoader(model_path=model_path) -model, tokenizer = ml.loader(num_gpus=1, load_8bit=ISLOAD_8BIT, debug=ISDEBUG) -#model, tokenizer = load_model(model_path=model_path, device=DEVICE, num_gpus=1, load_8bit=True, debug=False) class ModelWorker: - def __init__(self): - pass - # TODO + def __init__(self, model_path, model_name, device, num_gpus=1): + + if model_path.endswith("/"): + model_path = model_path[:-1] + self.model_name = model_name or model_path.split("/")[-1] + self.device = device + + self.ml = ModelLoader(model_path=model_path) + self.model, self.tokenizer = self.ml.loader(num_gpus, load_8bit=ISLOAD_8BIT, debug=ISDEBUG) + + if hasattr(self.model.config, "max_sequence_length"): + self.context_len = self.model.config.max_sequence_length + elif hasattr(self.model.config, "max_position_embeddings"): + self.context_len = self.model.config.max_position_embeddings + + else: + self.context_len = 2048 + + self.llm_chat_adapter = get_llm_chat_adapter(model_path) + self.generate_stream_func = self.llm_chat_adapter.get_generate_stream_func() + + def get_queue_length(self): + if model_semaphore is None or model_semaphore._value is None or model_semaphore._waiters is None: + return 0 + else: + CFG.LIMIT_MODEL_CONCURRENCY - model_semaphore._value + len(model_semaphore._waiters) + + def generate_stream_gate(self, params): + try: + for output in self.generate_stream_func( + self.model, + self.tokenizer, + params, + DEVICE, + CFG.MAX_POSITION_EMBEDDINGS + ): + print("output: ", output) + ret = { + "text": output, + "error_code": 0, + } + yield json.dumps(ret).encode() + b"\0" + + except torch.cuda.CudaError: + ret = { + "text": "**GPU OutOfMemory, Please Refresh.**", + "error_code": 0 + } + yield json.dumps(ret).encode() + b"\0" + + def get_embeddings(self, prompt): + return get_embeddings(self.model, self.tokenizer, prompt) app = FastAPI() @@ -61,41 +106,17 @@ def release_model_semaphore(): model_semaphore.release() -def generate_stream_gate(params): - try: - for output in generate_stream( - model, - tokenizer, - params, - DEVICE, - CFG.MAX_POSITION_EMBEDDINGS, - ): - print("output: ", output) - ret = { - "text": output, - "error_code": 0, - } - yield json.dumps(ret).encode() + b"\0" - except torch.cuda.CudaError: - ret = { - "text": "**GPU OutOfMemory, Please Refresh.**", - "error_code": 0 - } - yield json.dumps(ret).encode() + b"\0" - - @app.post("/generate_stream") async def api_generate_stream(request: Request): global model_semaphore, global_counter global_counter += 1 params = await request.json() - print(model, tokenizer, params, DEVICE) if model_semaphore is None: model_semaphore = asyncio.Semaphore(CFG.LIMIT_MODEL_CONCURRENCY) await model_semaphore.acquire() - generator = generate_stream_gate(params) + generator = worker.generate_stream_gate(params) background_tasks = BackgroundTasks() background_tasks.add_task(release_model_semaphore) return StreamingResponse(generator, background=background_tasks) @@ -111,7 +132,7 @@ def generate(prompt_request: PromptRequest): response = [] rsp_str = "" - output = generate_stream_gate(params) + output = worker.generate_stream_gate(params) for rsp in output: # rsp = rsp.decode("utf-8") rsp_str = str(rsp, "utf-8") @@ -125,9 +146,21 @@ def generate(prompt_request: PromptRequest): def embeddings(prompt_request: EmbeddingRequest): params = {"prompt": prompt_request.prompt} print("Received prompt: ", params["prompt"]) - output = get_embeddings(model, tokenizer, params["prompt"]) + output = worker.get_embeddings(params["prompt"]) return {"response": [float(x) for x in output]} if __name__ == "__main__": - uvicorn.run(app, host="0.0.0.0", log_level="info") \ No newline at end of file + + model_path = LLM_MODEL_CONFIG[CFG.LLM_MODEL] + print(model_path, DEVICE) + + + worker = ModelWorker( + model_path=model_path, + model_name=CFG.LLM_MODEL, + device=DEVICE, + num_gpus=1 + ) + + uvicorn.run(app, host="0.0.0.0", port=CFG.MODEL_PORT, log_level="info") \ No newline at end of file diff --git a/pilot/server/webserver.py b/pilot/server/webserver.py index bcf8f6385..270eff67f 100644 --- a/pilot/server/webserver.py +++ b/pilot/server/webserver.py @@ -369,8 +369,16 @@ def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, re for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"): if chunk: data = json.loads(chunk.decode()) + + """ TODO Multi mode output handler, rewrite this for multi model, use adapter mode. + """ if data["error_code"] == 0: - output = data["text"][skip_echo_len:].strip() + + if "vicuna" in CFG.LLM_MODEL: + output = data["text"][skip_echo_len:].strip() + else: + output = data["text"].strip() + output = post_process_code(output) state.messages[-1][-1] = output + "▌" yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5 @@ -445,7 +453,7 @@ def build_single_model_ui(): notice_markdown = """ # DB-GPT - [DB-GPT](https://github.com/csunny/DB-GPT) 是一个实验性的开源应用程序，它基于[FastChat](https://github.com/lm-sys/FastChat)，并使用vicuna-13b作为基础模型。此外，此程序结合了langchain和llama-index基于现有知识库进行In-Context Learning来对其进行数据库相关知识的增强。它可以进行SQL生成、SQL诊断、数据库知识问答等一系列的工作。 总的来说，它是一个用于数据库的复杂且创新的AI工具。如果您对如何在工作中使用或实施DB-GPT有任何具体问题，请联系我, 我会尽力提供帮助, 同时也欢迎大家参与到项目建设中, 做一些有趣的事情。 + [DB-GPT](https://github.com/csunny/DB-GPT) 是一个开源的以数据库为基础的GPT实验项目，使用本地化的GPT大模型与您的数据和环境进行交互，无数据泄露风险，100% 私密，100% 安全。 """ learn_more_markdown = """ ### Licence @@ -646,7 +654,6 @@ if __name__ == "__main__": cfg = Config() # dbs = get_database_list() - cfg.set_plugins(scan_plugins(cfg, cfg.debug_mode)) # 加载插件可执行命令 diff --git a/pilot/vector_store/file_loader.py b/pilot/vector_store/file_loader.py index 8703e2e4c..279d5343c 100644 --- a/pilot/vector_store/file_loader.py +++ b/pilot/vector_store/file_loader.py @@ -48,7 +48,7 @@ class KnownLedge2Vector: # vector_store.add_documents(documents=documents) else: documents = self.load_knownlege() - # reinit + # reinit vector_store = Chroma.from_documents(documents=documents, embedding=self.embeddings, persist_directory=persist_dir) diff --git a/pilot/vector_store/milvus_store.py b/pilot/vector_store/milvus_store.py index 6b06dcf00..5204e6b11 100644 --- a/pilot/vector_store/milvus_store.py +++ b/pilot/vector_store/milvus_store.py @@ -67,8 +67,8 @@ class MilvusStore(VectorStoreBase): def init_schema_and_load(self, vector_name, documents): """Create a Milvus collection, indexes it with HNSW, load document. Args: - documents (List[str]): Text to insert. vector_name (Embeddings): your collection name. + documents (List[str]): Text to insert. Returns: VectorStore: The MilvusStore vector store. """ @@ -203,21 +203,21 @@ class MilvusStore(VectorStoreBase): info = self.collection.describe() self.collection.load() - def insert(self, text, model_config) -> str: - """Add an embedding of data into milvus. - Args: - text (str): The raw text to construct embedding index. - Returns: - str: log. - """ - # embedding = get_ada_embedding(data) - embeddings = HuggingFaceEmbeddings(model_name=self.model_config["model_name"]) - result = self.collection.insert([embeddings.embed_documents(text), text]) - _text = ( - "Inserting data into memory at primary key: " - f"{result.primary_keys[0]}:\n data: {text}" - ) - return _text + # def insert(self, text, model_config) -> str: + # """Add an embedding of data into milvus. + # Args: + # text (str): The raw text to construct embedding index. + # Returns: + # str: log. + # """ + # # embedding = get_ada_embedding(data) + # embeddings = HuggingFaceEmbeddings(model_name=self.model_config["model_name"]) + # result = self.collection.insert([embeddings.embed_documents(text), text]) + # _text = ( + # "Inserting data into memory at primary key: " + # f"{result.primary_keys[0]}:\n data: {text}" + # ) + # return _text def _add_texts( self, diff --git a/requirements.txt b/requirements.txt index ba31d0d04..29e792451 100644 --- a/requirements.txt +++ b/requirements.txt @@ -42,6 +42,7 @@ tenacity==8.2.2 peft pycocoevalcap sentence-transformers +cpm_kernels umap-learn notebook gradio==3.23 @@ -61,6 +62,13 @@ langchain nltk python-dotenv==1.0.0 pymilvus +vcrpy +chromadb +markdown2 +colorama +playsound +distro +pypdf # Testing dependencies pytest @@ -70,11 +78,4 @@ pytest-benchmark pytest-cov pytest-integration pytest-mock -vcrpy -pytest-recording -chromadb -markdown2 -colorama -playsound -distro -pypdf \ No newline at end of file +pytest-recording \ No newline at end of file diff --git a/tools/knowlege_init.py b/tools/knowlege_init.py index fdc754e05..60010e4de 100644 --- a/tools/knowlege_init.py +++ b/tools/knowlege_init.py @@ -36,7 +36,7 @@ class LocalKnowledgeInit: if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--vector_name", type=str, default="keting") + parser.add_argument("--vector_name", type=str, default="default") parser.add_argument("--append", type=bool, default=False) parser.add_argument("--store_type", type=str, default="Chroma") args = parser.parse_args()