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

batch mode

Browse source
This commit is contained in:
Michael Yang 2023-06-28 14:33:10 -07:00
parent 9154298085
commit e63ad69388
1 changed files with 22 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
ollama/cmd/cli.py
View file

@ -54,15 +54,18 @@ def list_models(*args, **kwargs):
def generate(*args, **kwargs): def generate(*args, **kwargs):
if prompt := kwargs.get('prompt'): if prompt := kwargs.get('prompt'):
print('>>>', prompt, flush=True) print('>>>', prompt, flush=True)
print(flush=True)
generate_oneshot(*args, **kwargs) generate_oneshot(*args, **kwargs)
print(flush=True)
return return
return generate_interactive(*args, **kwargs) if sys.stdin.isatty():
return generate_interactive(*args, **kwargs)
return generate_batch(*args, **kwargs)
def generate_oneshot(*args, **kwargs): def generate_oneshot(*args, **kwargs):
print(flush=True)
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", [])
@ -70,20 +73,26 @@ def generate_oneshot(*args, **kwargs):
print(choices[0].get("text", ""), end="", flush=True) print(choices[0].get("text", ""), end="", flush=True)
# end with a new line # end with a new line
print() print(flush=True)
print(flush=True)
def generate_interactive(*args, **kwargs): def generate_interactive(*args, **kwargs):
print('>>> ', end='', flush=True) while True:
for line in sys.stdin:
if not sys.stdin.isatty():
print(line, end='')
print(flush=True)
kwargs.update({'prompt': line})
generate_oneshot(*args, **kwargs)
print(flush=True)
print('>>> ', end='', flush=True) print('>>> ', end='', flush=True)
line = next(sys.stdin)
if not line:
return
kwargs.update({"prompt": line})
generate_oneshot(*args, **kwargs)
def generate_batch(*args, **kwargs):
for line in sys.stdin:
print('>>> ', line, end='', flush=True)
kwargs.update({"prompt": line})
generate_oneshot(*args, **kwargs)
def add(model, models_home): def add(model, models_home):