baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

update generate to yield object

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

2
ollama/cmd/cli.py
View file

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

6
ollama/cmd/server.py
View file

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

2
ollama/engine.py
View file

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