Commit 972e91ef authored by Michael Yang's avatar Michael Yang
Browse files

update generate to yield object

parent 77eddba5
import os
import sys
import json
from pathlib import Path
from argparse import ArgumentParser
......@@ -67,7 +66,6 @@ def generate_oneshot(*args, **kwargs):
print(flush=True)
for output in engine.generate(*args, **kwargs):
output = json.loads(output)
choices = output.get("choices", [])
if len(choices) > 0:
print(choices[0].get("text", ""), end="", flush=True)
......
from aiohttp import web
import json
import aiohttp_cors
from aiohttp import web
from ollama import engine
......@@ -89,7 +90,8 @@ async def generate(request):
}
for output in engine.generate(model, prompt, **kwargs):
await response.write(output.encode("utf-8"))
output = json.dumps(output).encode('utf-8')
await response.write(output)
await response.write(b"\n")
return response
......@@ -33,7 +33,7 @@ def generate(model, prompt, models_home=".", llms={}, *args, **kwargs):
kwargs.update({"stream": True})
for output in llm(prompt, *args, **kwargs):
yield json.dumps(output)
yield output
def load(model, models_home=".", llms={}):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment