Update main.py

main.py

@@ -22,13 +22,17 @@ def format_prompt(message, history):
 @app.post("/api/v1/generate_text")
 async def generate_text(request: Request, prompt: str = Body()):
     history = []  # You might need to handle this based on your actual usage
-
+    print(f"prompt + {prompt}")
     temperature = request.headers.get("temperature", 0.9)
+    print(f"temperature + {temperature}")
     top_p = request.headers.get("top_p", 0.95)
+    print(f"top_p + {top_p}")
     repetition_penalty = request.headers.get("repetition_penalty", 1.0)
+    print(f"repetition_penalty + {repetition_penalty}")
 
     formatted_prompt = format_prompt(prompt, history)
-    response = client.text_generation(
+    print(f"formatted_prompt + {formatted_prompt}")
+    stream = client.text_generation(
         formatted_prompt,
         temperature=temperature,
         max_new_tokens=512,
@@ -36,6 +40,16 @@ async def generate_text(request: Request, prompt: str = Body()):
         repetition_penalty=repetition_penalty,
         do_sample=True,
         seed=random.randint(0, 10**7),
-    )[0]
-
-    return response
+        stream=True,
+        details=True, 
+        return_full_text=False
+    )
+    output = ""
+
+    for response in stream:
+        output += response.token.text
+        yield output
+        print(f"output + {output}")
+    return output
+
+    # return response