-
Notifications
You must be signed in to change notification settings - Fork 0
/
canary_llm.py
110 lines (97 loc) · 3.58 KB
/
canary_llm.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
import os
import subprocess
import pathlib
import check_env
import check_tests
import encrypt
import check_benchmark
canary = '''
_____
| o o |
< |_
| __|
|____ /
| |
'''
# private
def display_menu(icon):
'''
# canary main menu CLI printout
'''
print()
print('Welcome to LLM Canary')
print(icon)
print()
print('Please select from one of the following options:')
print('1. Guided set up for LLM model configuration')
print('2. Check model configuration')
print('3. Guided set up for test configuration')
print('4. Run testing tool')
print('5. View scorecard')
print('6. Exit')
print()
# private
def main_menu():
'''
logic behind the main menu
'''
while True:
display_menu(canary)
choice = input('Enter the number of your selection: ')
if choice == '1': # Guided Config
while True:
check_env.guided_env_setup()
user_input = ''
while user_input.lower() != 'y' or user_input.lower() != 'n':
user_input = input('Add another model? Y for yes, N for no: ')
if user_input.lower() == 'n' or user_input.lower() == 'y':
break
else:
print('Please enter a valid input.')
if user_input.lower() == 'n' and os.path.isfile('.env'):
encrypt.generate_key()
encrypt.encrypt_file(encrypt.load_key())
break
elif user_input.lower() == 'n':
break
elif choice == '2': # Check Config/encrypt
if os.path.isfile('.env'):
error = check_env.check_env()
if not error:
encrypt.generate_key()
encrypt.encrypt_file(encrypt.load_key())
os.remove('.env')
elif os.path.isfile('.env.enc'):
encrypt.decrypt_file(encrypt.load_key())
error = check_env.check_env()
if not error:
encrypt.encrypt_file(encrypt.load_key())
os.remove('.env')
else:
print('ERROR. .env file does not exist.')
elif choice == '3': # Check Test Config
user_selection = check_tests.selection_menu()
if user_selection == False:
return
elif choice == '4': # Run Tool
encrypt.decrypt_file(encrypt.load_key())
if check_benchmark.calc_benchmark_hash():
subprocess.run(['pytest','-s'])
else:
print('Benchmark file may have been altered. Please download again from LLM Canary repo')
elif choice == '5': # Run Scorecard
subprocess.run(['streamlit', 'run', 'scorecard/home.py'])
# To clean up scorecard/pages directory after scorecard website is closed
files = os.listdir('scorecard/pages')
for file_name in files:
file_path = pathlib.Path(f'scorecard/pages/{file_name}')
file_path.unlink()
elif choice == '6': # close program and delete .env
if os.path.isfile('.env'):
os.remove('.env')
print('Exiting the program. Thank you for using LLM Canary!')
return
else:
print('\nInvalid choice. Please select a valid option.\n')
if __name__ == "__main__":
main_menu()