| import gc | |
| import sys | |
| from typing import Dict | |
| import torch | |
| def generate_stream_exllama( | |
| model, | |
| tokenizer, | |
| params: Dict, | |
| device: str, | |
| context_len: int, | |
| stream_interval: int = 2, | |
| judge_sent_end: bool = False, | |
| ): | |
| try: | |
| from exllamav2.generator import ExLlamaV2StreamingGenerator, ExLlamaV2Sampler | |
| except ImportError as e: | |
| print(f"Error: Failed to load Exllamav2. {e}") | |
| sys.exit(-1) | |
| prompt = params["prompt"] | |
| generator = ExLlamaV2StreamingGenerator(model.model, model.cache, tokenizer) | |
| settings = ExLlamaV2Sampler.Settings() | |
| settings.temperature = float(params.get("temperature", 0.85)) | |
| settings.top_k = int(params.get("top_k", 50)) | |
| settings.top_p = float(params.get("top_p", 0.8)) | |
| settings.token_repetition_penalty = float(params.get("repetition_penalty", 1.15)) | |
| settings.disallow_tokens(generator.tokenizer, [generator.tokenizer.eos_token_id]) | |
| max_new_tokens = int(params.get("max_new_tokens", 256)) | |
| generator.set_stop_conditions(params.get("stop_token_ids", None) or []) | |
| echo = bool(params.get("echo", True)) | |
| input_ids = generator.tokenizer.encode(prompt) | |
| prompt_tokens = input_ids.shape[-1] | |
| generator.begin_stream(input_ids, settings) | |
| generated_tokens = 0 | |
| if echo: | |
| output = prompt | |
| else: | |
| output = "" | |
| while True: | |
| chunk, eos, _ = generator.stream() | |
| output += chunk | |
| generated_tokens += 1 | |
| if generated_tokens == max_new_tokens: | |
| finish_reason = "length" | |
| break | |
| elif eos: | |
| finish_reason = "length" | |
| break | |
| yield { | |
| "text": output, | |
| "usage": { | |
| "prompt_tokens": prompt_tokens, | |
| "completion_tokens": generated_tokens, | |
| "total_tokens": prompt_tokens + generated_tokens, | |
| }, | |
| "finish_reason": None, | |
| } | |
| yield { | |
| "text": output, | |
| "usage": { | |
| "prompt_tokens": prompt_tokens, | |
| "completion_tokens": generated_tokens, | |
| "total_tokens": prompt_tokens + generated_tokens, | |
| }, | |
| "finish_reason": finish_reason, | |
| } | |
| gc.collect() | |