yield

app.py

@@ -6,6 +6,15 @@ 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,
@@ -44,7 +53,7 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    # response = ""
+    response = ""
 
     # for message in client.chat_completion(
     #     messages,
@@ -65,14 +74,13 @@ def respond(
         return_tensors = "pt",
     )
     
-    # 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,
+    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)
     
-    # return tokenizer.batch_decode(outputs)
+    for token in custom_streamer.generated_text:
+        response += token
+        yield response
 
 
 """