test

app.py

@@ -6,15 +6,6 @@ For more information on `huggingface_hub` Inference API support, please check th
 """
 # client = InferenceClient("EITD/lora_model", token=os.getenv("HF_TOKEN"))
 
-class CustomTextStreamer(TextStreamer):
-    def __init__(self, tokenizer):
-        super().__init__(tokenizer)
-        self.generated_text = ""
-
-    def on_token(self, token):
-        super().on_token(token)
-        self.generated_text += token
-
 model = AutoPeftModelForCausalLM.from_pretrained(
         "EITD/lora_model_1", # YOUR MODEL YOU USED FOR TRAINING
         load_in_4bit = False,
@@ -53,7 +44,7 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    response = ""
+    # response = ""
 
     # for message in client.chat_completion(
     #     messages,
@@ -74,13 +65,14 @@ def respond(
         return_tensors = "pt",
     )
     
-    custom_streamer = CustomTextStreamer(tokenizer)
-    model.generate(input_ids = inputs, streamer = custom_streamer, max_new_tokens = max_tokens,
-                    use_cache = True, temperature = temperature, min_p = top_p)
+    outputs = model.generate(input_ids = inputs, max_new_tokens = max_tokens, use_cache = True,
+                         temperature = temperature, min_p = top_p)
+    
+    # text_streamer = TextStreamer(tokenizer, skip_prompt = True)
+    # model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = max_tokens,
+    #                 use_cache = True, temperature = temperature, min_p = top_p)
     
-    for token in custom_streamer.generated_text:
-        response += token
-        yield response
+    yield tokenizer.batch_decode(outputs, skip_special_tokens = True)
 
 
 """