(
model_path: str,
device: str,
num_gpus: str,
max_gpu_memory: str,
load_8bit: bool,
load_4bit: bool,
conv_template: Optional[str],
temperature: float,
max_new_tokens: int,
chatio: ChatIO,
debug: bool,
)
| 219 | |
| 220 | |
| 221 | def chat_loop( |
| 222 | model_path: str, |
| 223 | device: str, |
| 224 | num_gpus: str, |
| 225 | max_gpu_memory: str, |
| 226 | load_8bit: bool, |
| 227 | load_4bit: bool, |
| 228 | conv_template: Optional[str], |
| 229 | temperature: float, |
| 230 | max_new_tokens: int, |
| 231 | chatio: ChatIO, |
| 232 | debug: bool, |
| 233 | ): |
| 234 | # Model |
| 235 | model, tokenizer = load_model( |
| 236 | model_path, device, num_gpus, max_gpu_memory, load_8bit, load_4bit, debug |
| 237 | ) |
| 238 | |
| 239 | # Chat |
| 240 | if conv_template is None: |
| 241 | conv = get_default_conv_template().copy() |
| 242 | else: |
| 243 | raise NotImplementedError |
| 244 | |
| 245 | while True: |
| 246 | try: |
| 247 | inp = chatio.prompt_for_input(conv.roles[0]) |
| 248 | except EOFError: |
| 249 | inp = "" |
| 250 | if not inp: |
| 251 | print("exit...") |
| 252 | break |
| 253 | |
| 254 | conv.append_message(conv.roles[0], inp) |
| 255 | conv.append_message(conv.roles[1], None) |
| 256 | |
| 257 | generate_stream_func = generate_stream |
| 258 | prompt = conv.get_prompt() |
| 259 | |
| 260 | params = { |
| 261 | "model": model_path, |
| 262 | "prompt": prompt, |
| 263 | "temperature": temperature, |
| 264 | "max_new_tokens": max_new_tokens, |
| 265 | "stop": conv.sep if conv.sep_style == SeparatorStyle.SINGLE else None, |
| 266 | } |
| 267 | |
| 268 | chatio.prompt_for_output(conv.roles[1]) |
| 269 | output_stream = generate_stream_func(model, tokenizer, params, device) |
| 270 | outputs = chatio.stream_output(output_stream) |
| 271 | # NOTE: strip is important to align with the training data. |
| 272 | conv.messages[-1][-1] = outputs.strip() |
| 273 | |
| 274 | if debug: |
| 275 | print("\n", {"prompt": prompt, "outputs": outputs}, "\n") |
no test coverage detected