diff --git a/assistant.py b/assistant.py index cec98c4..da009a9 100755 --- a/assistant.py +++ b/assistant.py @@ -6,7 +6,8 @@ import pyperclip import sys import argparse -model = 'mistral-nemo' +model = 'llama3.1:8b-instruct-q8_0' +temp = 0.2 pattern = r'```[a-z]*\n[\s\S]*?\n```' line_pattern = r'`[a-z]*[\s\S]*?`' @@ -49,6 +50,7 @@ def chat(message, stream=True): history.append({"role": "user", "content": message}) completion = client.chat( model=model, + options={"temperature":temp}, messages=history, stream=stream ) @@ -56,7 +58,9 @@ def chat(message, stream=True): for chunk in completion: if stream: print(chunk['message']['content'], end='', flush=True) - result += chunk['message']['content'] + result += chunk['message']['content'] + if not stream: + result = completion['message']['content'] if stream: print() history.append({"role": 'assistant', 'content': result}) @@ -73,6 +77,8 @@ def parse_args(): parser.add_argument('--shell', '-s', nargs='?', const=True, default=False, help='output a shell command that does as described') # Add the --model (-m) argument parser.add_argument('--model', '-m', nargs='?', const=True, default=False, help='Specify model') + # Add the --temp (-t) argument + parser.add_argument('--temp', '-t', nargs='?', const=True, default=False, help='Specify temperature') # Parse the arguments return parser.parse_args() @@ -85,7 +91,7 @@ def arg_follow_up(args): return second_input def arg_shell(args): - query = 'Form a shell command based on the following description. Only output a working shell command and then <|eot_id|>.\nDescription: ' + query = 'Form a shell command based on the following description. Only output a working shell command .\nDescription: ' if args.shell != True: query += args.shell else: @@ -131,6 +137,9 @@ def main(): if args.model: global model model = args.model + if args.temp: + global temp + temp = float(args.temp) if not sys.stdin.isatty(): handle_piped_input(args) else: