Update main.py

main.py

@@ -68,11 +68,11 @@ async def generate_response(chat_chunks, llm):
                         'role': 'system',
                         'content': llm.detokenize(chat_chunk)
                     },
-                    'finish_reason': 'stop' if chat_chunk == "<|end|>" else 'unknown'
+                    'finish_reason': 'stop' if llm.is_eos_token(chat_chunk) else 'unknown'
                 }
             ]
         }
-        yield f"data: {json.dumps(response)}\n\n"
+        yield dict(data=json.dumps(response))
     yield dict(data="[DONE]")
 
 @app.post("/v1/chat/completions")
@@ -85,7 +85,7 @@ async def chat(request: ChatCompletionRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-    return StreamingResponse(generate_response(chat_chunks, llm), media_type="text/event-stream")
+    return EventSourceResponse(generate_response(chat_chunks, llm))
 
 async def stream_response(tokens: Any) -> None:
     try: