refactor

ollama/cmd/cli.py

@@ -0,0 +1,43 @@
+import json
+from pathlib import Path
+from argparse import ArgumentParser
+
+from ollama import model, engine
+from ollama.cmd import server
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument('--models-home', default=Path.home() / '.ollama' / 'models')
+
+    subparsers = parser.add_subparsers()
+
+    server.set_parser(subparsers.add_parser('serve'))
+
+    list_parser = subparsers.add_parser('list')
+    list_parser.set_defaults(fn=list)
+
+    generate_parser = subparsers.add_parser('generate')
+    generate_parser.add_argument('model')
+    generate_parser.add_argument('prompt')
+    generate_parser.set_defaults(fn=generate)
+
+    args = parser.parse_args()
+    args = vars(args)
+
+    fn = args.pop('fn')
+    fn(**args)
+
+
+def list(*args, **kwargs):
+    for m in model.models(*args, **kwargs):
+        print(m)
+
+
+def generate(*args, **kwargs):
+    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='')

ollama/cmd/server.py

@@ -0,0 +1,75 @@
+from aiohttp import web
+
+from ollama import engine
+
+
+def set_parser(parser):
+    parser.add_argument('--host', default='127.0.0.1')
+    parser.add_argument('--port', default=7734)
+    parser.set_defaults(fn=serve)
+
+
+def serve(models_home='.', *args, **kwargs):
+    app = web.Application()
+    app.add_routes([
+        web.post('/load', load),
+        web.post('/unload', unload),
+        web.post('/generate', generate),
+    ])
+
+    app.update({
+        'llms': {},
+        'models_home': models_home,
+    })
+
+    web.run_app(app, **kwargs)
+
+
+async def load(request):
+    body = await request.json()
+    model = body.get('model')
+    if not model:
+        raise web.HTTPBadRequest()
+
+    kwargs = {
+        'llms': request.app.get('llms'),
+        'models_home': request.app.get('models_home'),
+    }
+
+    engine.load(model, **kwargs)
+    return web.Response()
+
+
+async def unload(request):
+    body = await request.json()
+    model = body.get('model')
+    if not model:
+        raise web.HTTPBadRequest()
+
+    engine.unload(model, llms=request.app.get('llms'))
+    return web.Response()
+
+
+async def generate(request):
+    body = await request.json()
+    model = body.get('model')
+    if not model:
+        raise web.HTTPBadRequest()
+
+    prompt = body.get('prompt')
+    if not prompt:
+        raise web.HTTPBadRequest()
+
+    response = web.StreamResponse()
+    await response.prepare(request)
+
+    kwargs = {
+        'llms': request.app.get('llms'),
+        'models_home': request.app.get('models_home'),
+    }
+
+    for output in engine.generate(model, prompt, **kwargs):
+        await response.write(output.encode('utf-8'))
+        await response.write(b'\n')
+
+    return response