chewt/ai-assistant
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

feat: add temp flag

Browse source
This commit is contained in:
Hayden Johnson 2024-09-03 02:38:36 -07:00
parent 46e626b288
commit 65682cf598
1 changed files with 12 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
assistant.py
View file

@ -6,7 +6,8 @@ import pyperclip
import sys import sys
import argparse import argparse
model = 'mistral-nemo' model = 'llama3.1:8b-instruct-q8_0'
temp = 0.2
pattern = r'```[a-z]*\n[\s\S]*?\n```' pattern = r'```[a-z]*\n[\s\S]*?\n```'
line_pattern = r'`[a-z]*[\s\S]*?`' line_pattern = r'`[a-z]*[\s\S]*?`'
@ -49,6 +50,7 @@ def chat(message, stream=True):
history.append({"role": "user", "content": message}) history.append({"role": "user", "content": message})
completion = client.chat( completion = client.chat(
model=model, model=model,
options={"temperature":temp},
messages=history, messages=history,
stream=stream stream=stream
) )
@ -57,6 +59,8 @@ def chat(message, stream=True):
if stream: if stream:
print(chunk['message']['content'], end='', flush=True) print(chunk['message']['content'], end='', flush=True)
result += chunk['message']['content'] result += chunk['message']['content']
if not stream:
result = completion['message']['content']
if stream: if stream:
print() print()
history.append({"role": 'assistant', 'content': result}) 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') parser.add_argument('--shell', '-s', nargs='?', const=True, default=False, help='output a shell command that does as described')
# Add the --model (-m) argument # Add the --model (-m) argument
parser.add_argument('--model', '-m', nargs='?', const=True, default=False, help='Specify model') 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 # Parse the arguments
return parser.parse_args() return parser.parse_args()
@ -85,7 +91,7 @@ def arg_follow_up(args):
return second_input return second_input
def arg_shell(args): 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: if args.shell != True:
query += args.shell query += args.shell
else: else:
@ -131,6 +137,9 @@ def main():
if args.model: if args.model:
global model global model
model = args.model model = args.model
if args.temp:
global temp
temp = float(args.temp)
if not sys.stdin.isatty(): if not sys.stdin.isatty():
handle_piped_input(args) handle_piped_input(args)
else: else: