-
Notifications
You must be signed in to change notification settings - Fork 0
/
martha.py
executable file
·112 lines (93 loc) · 3.31 KB
/
martha.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
#!/usr/bin/env python3
from llama_cpp import Llama
from io import StringIO
import sys
import json
from pathlib import Path
import subprocess
CONVERSATION_FILE = Path("conversation_history.txt")
MODEL="./models/openhermes-2.5-mistral-7b-16k.Q5_K_M.gguf"
#MODEL="./models/Meta-Llama-3-8B-Instruct.Q5_K_M.gguf"
def run_python_code(code):
if not isinstance(code, str):
print(f'Found code "{code}" is not string...')
return
if code == "":
print(f'Found code "{code}" is empty string...')
return
out = ""
if input("Do you want to execute the code (y/n)? ") == "y":
# out = subprocess.run(['which', 'python'], capture_output=True, text=True)
# pyth_ver = out.stdout
# print(f"python: {pyth_ver}")
old_stdout = sys.stdout
with StringIO() as stringio:
redirected_output = sys.stdout = stringio
exec(code)
sys.stdout = old_stdout
out = f"Output:\n```\n{redirected_output.getvalue()}\n```\n"
return out
def main():
# initiate
assistant_name = Path(__file__).name
assistant_path = Path(__file__).absolute()
system_prompt = f"You are a python program living under {assistant_path}. {assistant_name} writes Python code to answer questions. When writing functions, she adds print statements in the code when helpful to understand the information flow. {assistant_name} keeps the code fairly short and doesn't read from user input. She provides single code snippet at a time."
# Mistral parameters
llm = Llama(
model_path=MODEL,
n_ctx=4096,
temperature=0.7,
verbose=False,
chat_format="chatml",
n_gpu_layers=1,
)
# Llama3 paremeters
# llm = Llama(
# model_path=MODEL,
# n_ctx=2048,
# temperature=0.2,
# verbose=False,
# chat_format="chatml",
# n_gpu_layers=1,
# )
if CONVERSATION_FILE.is_file():
with open(CONVERSATION_FILE, "r") as file:
messages = json.load(file)
else:
messages = [
{"role": "system", "content": system_prompt},
]
while True:
user_input = input("user: ")
messages.append(
{
"role": "user",
"content": user_input,
}
)
output_generator = llm.create_chat_completion(messages=messages)
role, content = (
output_generator["choices"][0]["message"]["role"],
output_generator["choices"][0]["message"]["content"],
)
print(f"{role}: {content}")
# handle python code execution
code_output = ""
if "```" in content and "```python\n" in content:
python_code = content.split("```python\n")[1].split("```")[0]
code_output = run_python_code(python_code)
print(code_output)
output_with_plugins = content + "\n" + code_output
#print(f"DEBUG: {output_with_plugins}")
# update conversation history
messages.append(
{
"role": "assistant",
"content": output_with_plugins,
}
)
json_dump = json.dumps(messages, indent=2)
with open(file=CONVERSATION_FILE, mode="w") as file:
file.write(json_dump)
if __name__ == "__main__":
main()