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():
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()
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)
generate_parser = subparsers.add_parser('generate')
generate_parser.add_argument('model')
generate_parser.add_argument('prompt')
generate_parser = subparsers.add_parser("generate")
generate_parser.add_argument("model")
generate_parser.add_argument("prompt")
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 = vars(args)
fn = args.pop('fn')
fn = args.pop("fn")
fn(**args)
@@ -38,6 +42,10 @@ def 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:
print(choices[0].get('text', ''), end='')
print(choices[0].get("text", ""), end="")
def add(*args, **kwargs):
model.add(*args, **kwargs)