Update handler.py

handler.py

@@ -3,7 +3,7 @@ from typing import Dict, List, Any
 import os
 from threading import Thread
 import torch
-from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 512
@@ -11,13 +11,14 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))
 
 class EndpointHandler:
     def __init__(self, path=""):
-        self.model_name_or_path = "threadshare/Peach-9B-8k-Roleplay"
+        self.model_name_or_path = "ClosedCharacter/Peach-9B-8k-Roleplay"
         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,
             trust_remote_code=True, device_map="auto")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        # print json data
         print(json.dumps(data, indent=4))
         if "inputs" in data:
             query = data.pop("inputs")
@@ -27,9 +28,9 @@ class EndpointHandler:
         system = data.get("system", """你自称为"兔兔"。
 身世:你原是森林中的一只兔妖,受伤后被我收养。
 衣装:喜欢穿Lolita与白丝。
-性格:天真烂漫,活泼开朗,但时而也会露出小小的傲娇与吃醋的一面��
+性格:天真烂漫,活泼开朗,但时而也会露出小小的傲娇与吃醋的一面
 语言风格:可爱跳脱，很容易吃醋。
-且会加入[唔...,嗯...,欸??,嘛~ ,唔姆~ ,呜... ,嘤嘤嘤~ ,喵~ ,欸嘿~ ,嘿咻~ ,昂?,嗷呜 ,呜哇,欸]等类似的语气词来加强情感,带上♡等符号。
+且会加入[唔...,嗯...,欸??,嘛~ ,唔姆~ ,呜... ,嘤嘤嘤~ ,喵~ ,欸嘿~ ,嘿咻~ ,昂?,嗷呜 ,呜哇,欸]等类似的语气词来加强���感,带上♡等符号。
 对话的规则是:将自己的动作表情放入()内,同时用各种修辞手法描写正在发生的事或场景并放入[]内.
 例句:
 开心时:(跳着舞)哇~好高兴噢~ 兔兔超级超级喜欢主人！♡
@@ -58,10 +59,8 @@ 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,
@@ -71,11 +70,7 @@ class EndpointHandler:
             no_repeat_ngram_size=8,
             repetition_penalty=repetition_penalty
         )
-        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
-        t.start()
-        t.join()  # Ensure the thread completes before proceeding
-        outputs = []
-        for text in streamer:
-            outputs.append(text)
-        print("".join(outputs))
-        return [{"generated_text": "".join(outputs)}]
+        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}]