This commit is contained in:
Bruce MacDonald 2023-06-27 17:26:31 -04:00
parent ece5259bc2
commit 7454900733
3 changed files with 62 additions and 36 deletions

View file

@ -8,24 +8,28 @@ from ollama.cmd import server
def main(): def main():
parser = ArgumentParser() parser = ArgumentParser()
parser.add_argument('--models-home', default=Path.home() / '.ollama' / 'models') parser.add_argument("--models-home", default=Path.home() / ".ollama" / "models")
subparsers = parser.add_subparsers() subparsers = parser.add_subparsers()
server.set_parser(subparsers.add_parser('serve')) server.set_parser(subparsers.add_parser("serve"))
list_parser = subparsers.add_parser('list') list_parser = subparsers.add_parser("list")
list_parser.set_defaults(fn=list) list_parser.set_defaults(fn=list)
generate_parser = subparsers.add_parser('generate') generate_parser = subparsers.add_parser("generate")
generate_parser.add_argument('model') generate_parser.add_argument("model")
generate_parser.add_argument('prompt') generate_parser.add_argument("prompt")
generate_parser.set_defaults(fn=generate) generate_parser.set_defaults(fn=generate)
add_parser = subparsers.add_parser("add")
add_parser.add_argument("model_file")
generate_parser.set_defaults(fn=add)
args = parser.parse_args() args = parser.parse_args()
args = vars(args) args = vars(args)
fn = args.pop('fn') fn = args.pop("fn")
fn(**args) fn(**args)
@ -38,6 +42,10 @@ def generate(*args, **kwargs):
for output in engine.generate(*args, **kwargs): for output in engine.generate(*args, **kwargs):
output = json.loads(output) 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='') print(choices[0].get("text", ""), end="")
def add(*args, **kwargs):
model.add(*args, **kwargs)

View file

@ -1,39 +1,59 @@
from aiohttp import web from aiohttp import web
import aiohttp_cors
from ollama import engine from ollama import engine
def set_parser(parser): def set_parser(parser):
parser.add_argument('--host', default='127.0.0.1') parser.add_argument("--host", default="127.0.0.1")
parser.add_argument('--port', default=7734) parser.add_argument("--port", default=7734)
parser.set_defaults(fn=serve) parser.set_defaults(fn=serve)
def serve(models_home='.', *args, **kwargs): def serve(models_home=".", *args, **kwargs):
app = web.Application() app = web.Application()
app.add_routes([
web.post('/load', load),
web.post('/unload', unload),
web.post('/generate', generate),
])
app.update({ cors = aiohttp_cors.setup(
'llms': {}, app,
'models_home': models_home, defaults={
}) "*": aiohttp_cors.ResourceOptions(
allow_credentials=True,
expose_headers="*",
allow_headers="*",
)
},
)
app.add_routes(
[
web.post("/load", load),
web.post("/unload", unload),
web.post("/generate", generate),
]
)
for route in app.router.routes():
cors.add(route)
app.update(
{
"llms": {},
"models_home": models_home,
}
)
web.run_app(app, **kwargs) web.run_app(app, **kwargs)
async def load(request): async def load(request):
body = await request.json() body = await request.json()
model = body.get('model') model = body.get("model")
if not model: if not model:
raise web.HTTPBadRequest() raise web.HTTPBadRequest()
kwargs = { kwargs = {
'llms': request.app.get('llms'), "llms": request.app.get("llms"),
'models_home': request.app.get('models_home'), "models_home": request.app.get("models_home"),
} }
engine.load(model, **kwargs) engine.load(model, **kwargs)
@ -42,21 +62,21 @@ async def load(request):
async def unload(request): async def unload(request):
body = await request.json() body = await request.json()
model = body.get('model') model = body.get("model")
if not model: if not model:
raise web.HTTPBadRequest() raise web.HTTPBadRequest()
engine.unload(model, llms=request.app.get('llms')) engine.unload(model, llms=request.app.get("llms"))
return web.Response() return web.Response()
async def generate(request): async def generate(request):
body = await request.json() body = await request.json()
model = body.get('model') model = body.get("model")
if not model: if not model:
raise web.HTTPBadRequest() raise web.HTTPBadRequest()
prompt = body.get('prompt') prompt = body.get("prompt")
if not prompt: if not prompt:
raise web.HTTPBadRequest() raise web.HTTPBadRequest()
@ -64,12 +84,12 @@ async def generate(request):
await response.prepare(request) await response.prepare(request)
kwargs = { kwargs = {
'llms': request.app.get('llms'), "llms": request.app.get("llms"),
'models_home': request.app.get('models_home'), "models_home": request.app.get("models_home"),
} }
for output in engine.generate(model, prompt, **kwargs): for output in engine.generate(model, prompt, **kwargs):
await response.write(output.encode('utf-8')) await response.write(output.encode("utf-8"))
await response.write(b'\n') await response.write(b"\n")
return response return response

View file

@ -1,7 +1,5 @@
click==8.1.3 aiohttp==3.8.4
Flask==2.3.2 aiohttp_cors==0.7.0
Flask_Cors==3.0.10
llama_cpp_python==0.1.65 llama_cpp_python==0.1.65
pyinstaller==5.13.0 pyinstaller==5.13.0
setuptools==65.6.3 setuptools==65.6.3
tqdm==4.65.0