diff --git a/README.md b/README.md index edee61a..0544a67 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,7 @@ The Web version's API for Android also supports other devices. See [Python sampl - [InternLM2/InternLM2.5](https://github.com/InternLM/InternLM) -- [Qwen/Qwen2](https://github.com/QwenLM/Qwen2) +- [Qwen1.5~2.5](https://github.com/QwenLM/Qwen2) - [puyu](https://internlm.openxlab.org.cn/) - [StepFun](https://platform.stepfun.com) - [KIMI](https://kimi.moonshot.cn) diff --git a/huixiangdou/service/llm_server_hybrid.py b/huixiangdou/service/llm_server_hybrid.py index 52c78d3..7f66ba6 100644 --- a/huixiangdou/service/llm_server_hybrid.py +++ b/huixiangdou/service/llm_server_hybrid.py @@ -17,9 +17,10 @@ from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from sse_starlette.sse import EventSourceResponse - +from transformers import TextIteratorStreamer import uvicorn from typing import List, Tuple +from threading import Thread def os_run(cmd: str): ret = os.popen(cmd) @@ -57,6 +58,7 @@ def __init__(self, model_path: str): self.model_path = model_path self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + self.streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True) model_path_lower = model_path.lower() @@ -69,14 +71,6 @@ def __init__(self, model_path: str): elif 'qwen1.5' in model_path_lower: self.model = AutoModelForCausalLM.from_pretrained( model_path, device_map='auto', trust_remote_code=True).eval() - elif 'qwen' in model_path_lower: - self.model = AutoModelForCausalLM.from_pretrained( - model_path, - device_map='auto', - trust_remote_code=True, - use_cache_quantization=True, - use_cache_kernel=True, - use_flash_attn=False).eval() elif 'internlm2_5' in model_path_lower: self.model = AutoModelForCausalLM.from_pretrained( model_path, @@ -112,17 +106,15 @@ async def chat_stream(self, prompt: str, history=[]): messages, tokenize=False, add_generation_prompt=True) model_inputs = self.tokenizer([text], return_tensors='pt').to('cuda') - generated_ids = self.model.generate(model_inputs.input_ids, - max_new_tokens=512, - top_k=1) - generated_ids = [ - output_ids[len(input_ids):] for input_ids, output_ids in zip( - model_inputs.input_ids, generated_ids) - ] - - output_text = self.tokenizer.batch_decode( - generated_ids, skip_special_tokens=True)[0] - yield output_text + + generation_kwargs = dict(model_inputs, streamer=self.streamer, max_new_tokens=512) + thread = Thread(target=self.model.generate, kwargs=generation_kwargs) + thread.start() + + for new_text in self.streamer: + yield new_text + + thread.join() elif type(self.model).__name__ == 'InternLM2ForCausalLM':