update generate to yield object

This commit is contained in:
Michael Yang 2023-06-28 18:01:18 -07:00
parent 77eddba5b3
commit 972e91ef46
3 changed files with 5 additions and 5 deletions

View file

@ -1,6 +1,5 @@
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)

View file

@ -1,5 +1,6 @@
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

View file

@ -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={}):