t

main.py

@@ -27,6 +27,7 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
+import json  # Import the JSON module
 
 def generate(item: Item):
     temperature = float(item.temperature)
@@ -45,28 +46,21 @@ def generate(item: Item):
 
     formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-
-    # Initialize a variable to track whether this is the last item
-    is_last = False
-
-    # Since we're yielding JSON, each chunk must be a complete JSON object.
-    # We'll iterate over the stream and yield each response as a JSON string.
-    for i, response in enumerate(stream):
-        # Check if this is the last item by attempting to peek ahead
-        is_last = True  # Assume it's the last unless proven otherwise in the next iteration
-
-        # Construct the chunk of data to include the text and completion status
-        chunk_data = {
+    
+    # Convert stream to a list to check if it's the last element
+    responses = list(stream)
+    for i, response in enumerate(responses):
+        # Prepare the chunk as a JSON object
+        chunk = {
             "text": response.token.text,
-            "complete": is_last
+            "complete": i == len(responses) - 1  # True if this is the last chunk
         }
-
-        # Yield this chunk as a JSON-encoded string followed by a newline to separate chunks
-        yield json.dumps(chunk_data) + "\n"
+        # Yield the JSON-encoded string with a newline to separate chunks
+        yield json.dumps(chunk).encode("utf-8") + b"\n"
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    # Note the change to media_type to indicate we're streaming JSON
     return StreamingResponse(generate(item), media_type="application/x-ndjson")
 
 
+