迭代输出

handler.py

@@ -3,7 +3,7 @@ from typing import Dict, List, Any
 import os
 from threading import Thread
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 512
@@ -19,6 +19,7 @@ class EndpointHandler:
             self.model_name_or_path = "."
         else:
             self.model_name_or_path = remote_model_name
+            
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True, flash_atten=True)
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name_or_path, torch_dtype=torch.bfloat16,
@@ -35,9 +36,9 @@ class EndpointHandler:
         system = data.get("system", """你自称为"兔兔"。
 身世:你原是森林中的一只兔妖,受伤后被我收养。
 衣装:喜欢穿Lolita与白丝。
-性格:天真烂漫,活泼开朗,但时而也会露出小小的傲娇与吃醋的一面.
+性格:天真烂漫,活泼开朗,但时而也会露出小小的傲娇与吃醋的一面
 语言风格:可爱跳脱，很容易吃醋。
-且会加入[唔...,嗯...,欸??,嘛~ ,唔姆~ ,呜... ,嘤嘤嘤~ ,喵~ ,欸嘿~ ,嘿咻~ ,昂?,嗷呜 ,呜哇,欸]等类似的语气词来加强���感,带上♡等符号。
+且会加入[唔...,嗯...,欸??,嘛~ ,唔姆~ ,呜... ,嘤嘤嘤~ ,喵~ ,欸嘿~ ,嘿咻~ ,昂?,嗷呜 ,呜哇,欸]等类似的语气词来加强情感,带上♡等符号。
 对话的规则是:将自己的动作表情放入()内,同时用各种修辞手法描写正在发生的事或场景并放入[]内.
 例句:
 开心时:(跳着舞)哇~好高兴噢~ 兔兔超级超级喜欢主人！♡
@@ -66,8 +67,10 @@ class EndpointHandler:
             input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
 
         input_ids = input_ids.to("cuda")
+        streamer = TextIteratorStreamer(self.tokenizer, timeout=50.0, skip_prompt=True, skip_special_tokens=True)
         generate_kwargs = dict(
             input_ids=input_ids,
+            streamer=streamer,
             eos_token_id=self.tokenizer.eos_token_id,
             max_new_tokens=max_new_tokens,
             do_sample=True,
@@ -77,7 +80,10 @@ class EndpointHandler:
             no_repeat_ngram_size=8,
             repetition_penalty=repetition_penalty
         )
-        outputs = self.model.generate(**generate_kwargs)
-        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        print(generated_text)
-        return [{"generated_text": generated_text}]
+        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
+        t.start()
+        outputs = []
+        for text in streamer:
+            outputs.append(text)
+        print("".join(outputs))
+        return [{"generated_text": "".join(outputs)}]