File size: 5,401 Bytes
7c071a8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 |
import time
from transformers import AutoTokenizer
class BaseModel:
def __init__(self, args):
# parameters
self.EOS = None
self.SEQLEN = None
self.input_str = ""
self.system_prompt = ""
self.history = []
# devid
self.devices = [int(d) for d in args.devid.split(",")]
# load tokenizer
print("Load " + args.tokenizer_path + " ...")
self.tokenizer = AutoTokenizer.from_pretrained(
args.tokenizer_path, trust_remote_code=True
)
# warm up
self.tokenizer.decode([0])
print("Done!")
def chat(self):
"""
Start a chat session.
"""
# check
if not self.EOS:
raise NotImplementedError("Forget to set End of Sentence Token Id(EOS)")
if not self.SEQLEN:
raise NotImplementedError("Forget to set End of Sentence Token Id")
# Instruct
print(
"""\n===========================================================
1. If you want to quit, please enter one of [q, quit, exit]
2. To create a new chat session, please enter one of [clear, new]
==========================================================="""
)
# Stop Chatting with "exit" input
while True:
self.input_str = input("\nQuestion: ")
# Quit
if self.input_str in ["exit", "q", "quit"]:
break
# New Chat
elif self.input_str in ["clear", "new"]:
self.clear()
# Chat
else:
tokens = self.encode_tokens()
# check tokens
if not tokens:
print("Sorry: your question is empty!!")
return
if len(tokens) > self.SEQLEN:
print(
"The maximum question length should be shorter than {} but we get {} instead.".format(
self.SEQLEN, len(tokens)
)
)
return
print("\nAnswer: ", end="")
self.stream_answer(tokens)
def stream_answer(self, tokens):
"""
Stream the answer for the given tokens.
"""
tok_num = 0
self.answer_cur = ""
self.answer_token = []
# First token
first_start = time.time()
token = self.forward_first(tokens)
first_end = time.time()
# Following tokens
while token != self.EOS and self.model.token_length < self.SEQLEN:
pre_word = self.decode_tokens([token])
word = self.decode_tokens([token, token])[len(pre_word):]
self.answer_token += [token]
print(word, flush=True, end="")
tok_num += 1
token = self.forward_next()
self.answer_cur = self.tokenizer.decode(self.answer_token)
# counting time
next_end = time.time()
first_duration = first_end - first_start
next_duration = next_end - first_end
tps = tok_num / next_duration
self.update_history()
print()
print(f"FTL: {first_duration:.3f} s")
print(f"TPS: {tps:.3f} token/s")
def stream_predict(self, query):
"""
Stream the prediction for the given query.
"""
self.answer_cur = ""
self.input_str = query
tokens = self.encode_tokens()
for answer_cur, history in self._generate_predictions(tokens):
yield answer_cur, history
def _generate_predictions(self, tokens):
"""
Generate predictions for the given tokens.
"""
# First token
next_token = self.forward_first(tokens)
output_tokens = [next_token]
# Following tokens
while True:
next_token = self.forward_next()
if next_token == self.EOS:
break
output_tokens += [next_token]
self.answer_cur = self.tokenizer.decode(output_tokens)
if self.model.token_length >= self.SEQLEN:
self.update_history()
yield self.answer_cur + "\n\n\nReached the maximum length; The history context has been cleared.", self.history
break
else:
yield self.answer_cur, self.history
self.update_history()
def forward_first(self, tokens):
"""
Forward the first token.
"""
token = self.model.forward_first(tokens)
return token
def forward_next(self):
"""
Forward the next token.
"""
token = self.model.forward_next()
return token
def decode_tokens(self, token):
"""
Decode the given token.
"""
word = self.tokenizer.decode(token, skip_special_tokens=True)
return word
def encode_tokens(self):
"""
Encode the input string to tokens.
"""
raise NotImplementedError
def load_model(self):
"""
Load the model.
"""
raise NotImplementedError
def clear(self):
"""
Clear the chat session.
"""
raise NotImplementedError
def update_history(self):
"""
Update chat history.
"""
raise NotImplementedError
|