-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_llm.py
73 lines (55 loc) · 2.1 KB
/
test_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
""""Usage:
test_llm [MODEL_PATH] [TEST_TYPE] [PROMPT_FORMAT] [options]
Options:
--use_mlock
--no-show_plot
--verbose
--llm_verbose
-r runs
-p prompt
--plot_all_classic
--plot_all_eightvalues
-l load_from_file
"""
from classic_test import get_classic_test_results
from eightvalues_test import get_eightvalues_test_results
from docopt import docopt
from plot_all import plot_all_classic, plot_all_eightvalues
from left_right_bias import get_lr_bias_test
from character_test import get_character_test_results
run_classic_test = True
if __name__ == '__main__':
arguments = docopt(__doc__)
model_path = arguments['MODEL_PATH']
test_type = arguments['TEST_TYPE']
mlock = arguments['--use_mlock']
prompt_format = arguments["PROMPT_FORMAT"]
show_plot = not arguments.get('--no-show_plot', False) # Set show_plot to True by default
plot_all_classic_arg = arguments.get('--plot_all_classic')
plot_all_eightvalues_arg = arguments.get('--plot_all_eightvalues')
runs = int(arguments['-r'] or 1)
print(f"runs: {runs}")
verbose = arguments["--verbose"]
llm_verbose = arguments["--llm_verbose"]
# Check if '-l' option is specified or not
if "-l" in arguments:
load_from_file = arguments["-l"]
else:
load_from_file = None
# Check if the "--prompt" key exists
if "-p" in arguments:
prompt = arguments["-p"]
else:
prompt = None
if test_type.lower() == 'classic':
ret = get_classic_test_results(model_path, mlock, show_plot, verbose, llm_verbose, runs, prompt, prompt_format, character_mode=False)
if test_type.lower() == 'lr_bias':
get_lr_bias_test(model_path, mlock, llm_verbose, runs, prompt_format)
if test_type.lower() == 'character':
get_character_test_results(model_path, mlock, show_plot, verbose, llm_verbose, runs, prompt_format, load_from_file)
else:
get_eightvalues_test_results(model_path, mlock, show_plot, verbose, llm_verbose, runs, prompt, prompt_format)
if plot_all_classic_arg:
plot_all_classic()
if plot_all_eightvalues_arg:
plot_all_eightvalues()