Fix a bug in generate_text

app.py

@@ -34,12 +34,10 @@ def generate_text(prompt: str,
                   temperature: float = 0.5,
                   top_p: float = 0.95,
                   top_k: int = 50) -> str:
-
     # Encode the prompt
     inputs = tokenizer([prompt],
                        return_tensors='pt',
                        add_special_tokens=False).to(DEVICE)
-
     # Prepare arguments for generation
     input_length = inputs["input_ids"].shape[-1]
     max_new_tokens = min(max_new_tokens, WINDOW_SIZE - input_length)
@@ -56,8 +54,8 @@ def generate_text(prompt: str,
                                     skip_prompt=True,
                                     skip_special_tokens=True)
     generation_kwargs = dict(
-        inputs=inputs,
-        streamer=inputs,
+        **inputs,
+        streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
         top_p=top_p,
@@ -65,12 +63,10 @@ def generate_text(prompt: str,
         temperature=temperature,
         num_beams=1,
     )
-
     # Generate text
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-
-    generated_text = ""
+    generated_text = prompt
     for new_text in streamer:
         generated_text += new_text
     return generated_text