| import os |
| from typing import Dict, List, Optional, Tuple, Union |
|
|
| import torch |
| from fastchat.conversation import (compute_skip_echo_len, |
| get_default_conv_template) |
| from fastchat.serve.inference import load_model as load_fastchat_model |
| from langchain.llms.base import LLM |
| from langchain.llms.utils import enforce_stop_tokens |
| from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
| MODEL_CACHE_PATH = os.path.join(os.path.dirname(__file__), 'model_cache') |
|
|
| llm_model_dict = { |
| "chatglm": { |
| "ChatGLM-6B": "THUDM/chatglm-6b", |
| "ChatGLM-6B-int4": "THUDM/chatglm-6b-int4", |
| "ChatGLM-6B-int8": "THUDM/chatglm-6b-int8", |
| "ChatGLM-6b-int4-qe": "THUDM/chatglm-6b-int4-qe" |
| }, |
| "belle": { |
| "BELLE-LLaMA-Local": "/pretrainmodel/belle", |
| }, |
| "vicuna": { |
| "Vicuna-Local": "/pretrainmodel/vicuna", |
| } |
| } |
|
|
| os.environ["TOKENIZERS_PARALLELISM"] = "false" |
|
|
| DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" |
| DEVICE_ID = "0" |
| CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE |
|
|
|
|
| def torch_gc(): |
| if torch.cuda.is_available(): |
| with torch.cuda.device(CUDA_DEVICE): |
| torch.cuda.empty_cache() |
| torch.cuda.ipc_collect() |
|
|
|
|
| def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: |
| num_trans_layers = 28 |
| per_gpu_layers = 30 / num_gpus |
|
|
| device_map = { |
| 'transformer.word_embeddings': 0, |
| 'transformer.final_layernorm': 0, |
| 'lm_head': 0 |
| } |
|
|
| used = 2 |
| gpu_target = 0 |
| for i in range(num_trans_layers): |
| if used >= per_gpu_layers: |
| gpu_target += 1 |
| used = 0 |
| assert gpu_target < num_gpus |
| device_map[f'transformer.layers.{i}'] = gpu_target |
| used += 1 |
|
|
| return device_map |
|
|
|
|
| class ChatLLM(LLM): |
| max_token: int = 10000 |
| temperature: float = 0.1 |
| top_p = 0.9 |
| history = [] |
| model_type: str = "chatglm" |
| model_name_or_path: str = "ChatGLM-6B-int4", |
| tokenizer: object = None |
| model: object = None |
|
|
| def __init__(self): |
| super().__init__() |
|
|
| @property |
| def _llm_type(self) -> str: |
| return "ChatLLM" |
|
|
| def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: |
|
|
| if self.model_type == 'vicuna': |
| conv = get_default_conv_template(self.model_name_or_path).copy() |
| conv.append_message(conv.roles[0], prompt) |
| conv.append_message(conv.roles[1], None) |
| prompt = conv.get_prompt() |
| inputs = self.tokenizer([prompt]) |
| output_ids = self.model.generate( |
| torch.as_tensor(inputs.input_ids).cuda(), |
| do_sample=True, |
| temperature=self.temperature, |
| max_new_tokens=self.max_token, |
| ) |
| outputs = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0] |
| skip_echo_len = compute_skip_echo_len(self.model_name_or_path, conv, prompt) |
| response = outputs[skip_echo_len:] |
| torch_gc() |
| if stop is not None: |
| response = enforce_stop_tokens(response, stop) |
| self.history = [[None, response]] |
|
|
| elif self.model_type == 'belle': |
| prompt = "Human: "+ prompt +" \n\nAssistant: " |
| input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(DEVICE) |
| generate_ids = self.model.generate(input_ids, max_new_tokens=self.max_token, do_sample = True, top_k = 30, top_p = self.top_p, temperature = self.temperature, repetition_penalty=1., eos_token_id=2, bos_token_id=1, pad_token_id=0) |
| output = self.tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] |
| response = output[len(prompt)+1:] |
| torch_gc() |
| if stop is not None: |
| response = enforce_stop_tokens(response, stop) |
| self.history = [[None, response]] |
|
|
| elif self.model_type == 'chatglm': |
| response, _ = self.model.chat( |
| self.tokenizer, |
| prompt, |
| history=self.history, |
| max_length=self.max_token, |
| temperature=self.temperature, |
| ) |
| torch_gc() |
| if stop is not None: |
| response = enforce_stop_tokens(response, stop) |
| self.history = self.history + [[None, response]] |
|
|
| return response |
|
|
|
|
| def load_llm(self, |
| llm_device=DEVICE, |
| num_gpus=torch.cuda.device_count(), |
| device_map: Optional[Dict[str, int]] = None, |
| **kwargs): |
| if 'chatglm' in self.model_name_or_path.lower(): |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, |
| trust_remote_code=True, cache_dir=os.path.join(MODEL_CACHE_PATH, self.model_name_or_path)) |
| if torch.cuda.is_available() and llm_device.lower().startswith("cuda"): |
|
|
| num_gpus = torch.cuda.device_count() |
| if num_gpus < 2 and device_map is None: |
| self.model = (AutoModel.from_pretrained( |
| self.model_name_or_path, trust_remote_code=True, cache_dir=os.path.join(MODEL_CACHE_PATH, self.model_name_or_path), |
| **kwargs).half().cuda()) |
| else: |
| from accelerate import dispatch_model |
|
|
| model = AutoModel.from_pretrained(self.model_name_or_path, |
| trust_remote_code=True, cache_dir=os.path.join(MODEL_CACHE_PATH, self.model_name_or_path), |
| **kwargs).half() |
|
|
| if device_map is None: |
| device_map = auto_configure_device_map(num_gpus) |
|
|
| self.model = dispatch_model(model, device_map=device_map) |
| else: |
| self.model = (AutoModel.from_pretrained( |
| self.model_name_or_path, |
| trust_remote_code=True, cache_dir=os.path.join(MODEL_CACHE_PATH, self.model_name_or_path)).float().to(llm_device)) |
| self.model = self.model.eval() |
|
|
| else: |
| self.model, self.tokenizer = load_fastchat_model( |
| model_path = self.model_name_or_path, |
| device = llm_device, |
| num_gpus = num_gpus |
| ) |
| |
