|
|
from typing import List, Optional |
|
|
from langchain.llms.base import LLM |
|
|
import torch |
|
|
from transformers import AutoModel, AutoTokenizer |
|
|
from langchain.llms.utils import enforce_stop_tokens |
|
|
from fastchat.conversation import (compute_skip_echo_len, |
|
|
get_default_conv_template) |
|
|
|
|
|
|
|
|
class ModelLoader(LLM): |
|
|
tokenizer: object = None |
|
|
model: object = None |
|
|
max_token: int = 10000 |
|
|
temperature: float = 0.1 |
|
|
top_p = 0.9 |
|
|
history = [] |
|
|
|
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
|
|
|
@property |
|
|
def _llm_type(self) -> str: |
|
|
return "ChatLLM" |
|
|
|
|
|
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: |
|
|
conv = get_default_conv_template("/DATA/gpt/lang/model_cache/THUDM/chatglm-6b-int8").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("/DATA/gpt/lang/model_cache/THUDM/chatglm-6b-int8", conv, prompt) |
|
|
response = outputs[skip_echo_len:] |
|
|
if stop is not None: |
|
|
response = enforce_stop_tokens(response, stop) |
|
|
self.history = [[None, response]] |
|
|
|
|
|
return response |
|
|
|
|
|
|
|
|
def load_model(self, model_name_or_path: str = "/DATA/gpt/lang/model_cache/THUDM/chatglm-6b-int8"): |
|
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained( |
|
|
"/DATA/gpt/mingpt-7b/MiniGPT-4-LLaMA-7B", |
|
|
trust_remote_code=True |
|
|
) |
|
|
self.model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True) |
|
|
self.model = self.model.eval() |
|
|
|
