Unverified Commit a8d22fe4 authored by Michael Yang's avatar Michael Yang Committed by GitHub
Browse files

Merge pull request #18 from jmorganca/yield-object

update generate to yield object
parents 4819aa47 972e91ef
import os import os
import sys import sys
import json
from pathlib import Path from pathlib import Path
from argparse import ArgumentParser from argparse import ArgumentParser
...@@ -71,7 +70,6 @@ def generate_oneshot(*args, **kwargs): ...@@ -71,7 +70,6 @@ def generate_oneshot(*args, **kwargs):
print(flush=True) print(flush=True)
for output in engine.generate(*args, **kwargs): for output in engine.generate(*args, **kwargs):
output = json.loads(output)
choices = output.get("choices", []) choices = output.get("choices", [])
if len(choices) > 0: if len(choices) > 0:
print(choices[0].get("text", ""), end="", flush=True) print(choices[0].get("text", ""), end="", flush=True)
......
from aiohttp import web import json
import aiohttp_cors import aiohttp_cors
from aiohttp import web
from ollama import engine from ollama import engine
...@@ -89,7 +90,8 @@ async def generate(request): ...@@ -89,7 +90,8 @@ async def generate(request):
} }
for output in engine.generate(model, prompt, **kwargs): 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") await response.write(b"\n")
return response return response
...@@ -33,7 +33,7 @@ def generate(model, prompt, models_home=".", llms={}, *args, **kwargs): ...@@ -33,7 +33,7 @@ def generate(model, prompt, models_home=".", llms={}, *args, **kwargs):
kwargs.update({"stream": True}) kwargs.update({"stream": True})
for output in llm(prompt, *args, **kwargs): for output in llm(prompt, *args, **kwargs):
yield json.dumps(output) yield output
def load(model, models_home=".", llms={}): 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