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':
|