File size: 9,857 Bytes
ee7776a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e3508f0
 
ee7776a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0613fa7
ee7776a
 
 
 
 
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
import gradio as gr
import os
import run_auditor_user_defined
import run_critic_user_defined
import run_rank
import pre_process
from utils import dotdict, clean_folder

os.environ['DISPLAY'] = ':0'

def openAIKey(key):
    if key:
        os.environ["OPENAI_API_KEY"] = key
        return "Let's get you started"
    else:
        return "Please add your OpenAI API key to continue."


def auditor(key, audit_model, audit_top_k, audit_temperature, audit_number, audit_contract_file, audit_prompt_file):
    if audit_contract_file:
        args = {
            'backend': audit_model.lower(),
            'temperature': audit_temperature,
            'data_dir': "data/CVE_clean",
            'topk': audit_top_k,
            'num_auditor': audit_number,
            'openai_api_key': key
        }
        args = dotdict(args)

        if os.path.exists("data/CVE"):
            clean_folder("data/CVE")
        if os.path.exists("data/CVE_clean"):
            clean_folder("data/CVE_clean")
        if os.path.exists(f"src/logs/auditor_{args.backend}_{args.temperature}_top{args.topk}_{args.num_auditor}"):
            clean_folder(f"src/logs/auditor_{args.backend}_{args.temperature}_top{args.topk}_{args.num_auditor}")

        for uploaded_file in audit_contract_file:
            file_name_int = uploaded_file.rfind("/") + 1
            name = uploaded_file[file_name_int:len(uploaded_file)]
            fp = open(uploaded_file, 'rb')
            bytes_data = fp.read()
            with open(f"data/CVE/{name}", "wb") as f:
                f.write(bytes_data)
        
        ## Pre processing the data        
        pre_process.mainfnc(args.data_dir)

        if audit_prompt_file:
            for uploaded_file in audit_prompt_file:
                file_name_int = uploaded_file.rfind("/") + 1
                name = uploaded_file[file_name_int:len(uploaded_file)]
                fp = open(uploaded_file, 'rb')
                bytes_data = fp.read()
                with open(f"src/prompt.py", "wb") as f:
                    f.write(bytes_data)
        
        ## Running the auditor code
        run_auditor_user_defined.mainfnc(args)

        message = f"You can find your processed Audit files at ./src/logs/auditor_{args.backend}_{args.temperature}_top{args.topk}_{args.num_auditor}"
    else:
        message = "Please Upload the Contract Files to proceed"
    
    return message



def critic(key, critic_model, critic_shots, critic_temperature, critic_number, critic_auit_directory):
    if critic_auit_directory:
        args = {
            'backend': critic_model.lower(),
            'temperature': critic_temperature,
            'dataset': "CVE",
            'auditor_dir': critic_auit_directory,
            'num_critic': critic_number,
            'shot': critic_shots.lower(),
            'openai_api_key': key
        }
        args = dotdict(args)

        run_critic_user_defined.mainfnc(args)
        message = f"Critic files processed successfully to folder ./src/logs/{args.auditor_dir}/critic_{args.backend}_{args.temperature}_{args.num_critic}_{args.shot}"
    else:
        message = "Please add the Audit Directory value to proceed"

    return message


def ranker(critic_model, critic_shots, critic_temperature, critic_number, critic_auit_directory, ranker_strategy, ranker_audit_location, ranker_critic_location):
    if critic_auit_directory and ranker_audit_location:
        args_c = {
            'auditor_dir': critic_auit_directory,
            'backend':critic_model.lower(),
            'temperature':critic_temperature,
            'num_critic': critic_number,
            'shot': critic_shots.lower()
        }

        args_r = {
            'auditor_dir': ranker_audit_location,
            'critic_dir': ranker_critic_location,
            'strategy': ranker_strategy.lower()
        }
        args_r = dotdict(args_r)
        args_c = dotdict(args_c)

        run_rank.mainfnc(args_r)
        message = f"Ranking files processed successfully to folder ./src/logs/{args_c.auditor_dir}/critic_{args_c.backend}_{args_c.temperature}_{args_c.num_critic}_{args_c.shot}/ranker_{args_r.strategy}"
    else:
        message = "Please enter the Audit directory locations to proceed."

    return message



with gr.Blocks(theme=gr.themes.Base) as ui:
    gr.Image("src/images/header.jpg", min_width=1500, height=80, show_download_button=False, show_label=False)
    message = "Hello, Welcome to GPTLens GUI by Aishwarya Solanki..."
    audit_output = "Auditor hasn't run yet."
    critic_output = "Critic hasn't run yet."
    ranker_output = "Ranker hasn't run yet."

    with gr.Row(equal_height=500):
        with gr.Column(scale=1, min_width=300):
            key = gr.Textbox(label="Please input your openAi key", type="password")
            gr.Button(value="Don't have an openAi key? Get one here.", link="https://platform.openai.com/account/api-keys")
            gr.ClearButton(key, message, variant="primary")
            

        with gr.Column(scale=2, min_width=600):
            message = gr.Textbox(label="", value=message, interactive=False, visible=True)
            gr.Markdown("### Auditor")

            audit_model = gr.Radio(label="GPT Model", value="GPT-4",  choices=["GPT-3.5-Turbo", "GPT-4", "GPT-4-Turbo-Preview"])
            audit_top_k = gr.Slider(1, 10, value=3, step = 1, label="top-k value", info="Set the 'k' for top-k auditor responses")
            audit_temperature = gr.Slider(0.0, 1.0, value=0.7, step = 0.01, label="Temperature", info="Set the Temperature")
            audit_number = gr.Slider(1, 10, value=1, step = 1, label="Number of Auditors", info="Set the number of auditors")

            with gr.Row():
                with gr.Column(min_width=300, scale=2):
                    gr.Markdown("### Upload smart contract files (.sol)")
                    audit_contract_file = gr.Files()
                    audit_btn = gr.Button(value="Start Auditor", variant="primary", elem_id="audit")

                with gr.Column(min_width=300, scale=2):
                    gr.Markdown("### Upload prompt file (optional, .py)")
                    audit_prompt_file = gr.Files()
            
            audit_output = gr.Textbox(label="Auditor Status", value=audit_output, interactive=False, visible=True, elem_id="auditor_output")
            
            audit_btn.click(fn=auditor, inputs=[key, audit_model, audit_top_k, audit_temperature, 
                                                audit_number, audit_contract_file, audit_prompt_file], 
                                                outputs=audit_output, api_name="auditor")
            gr.Markdown("### Critic")
            critic_model = gr.Radio(label="GPT Model", value="GPT-4", choices=["GPT-3.5-Turbo", "GPT-4", "GPT-4-Turbo-Preview"])
            critic_shots= gr.Radio(label="Set the num of shots", value="Few", choices=["One", "Few"])
            critic_temperature = gr.Slider(0.0, 1.0, value=0.7, step = 0.01, label="Temperature", info="Set the Temperature")
            critic_number = gr.Slider(1, 10, value=1, step = 1, label="Number of Critics", info="Set the number of critics")
            critic_auit_directory = gr.Textbox(label="Auditor Directory Location", 
                                               info="Example = auditor_{audit_model}_{audit_temperature}_top{audit_top_k}_{audit_number}", 
                                               value="")

            with gr.Row():
                with gr.Column(min_width=300, scale=2):
                    critic_btn = gr.Button(value="Start Critic", variant="primary", elem_id="critic")
                with gr.Column(min_width=300, scale=2):
                    gr.Button(value="Start Ranker", variant="primary", visible=False)

            critic_output = gr.Textbox(label="Critic Status", value=critic_output, interactive=False, visible=True, elem_id="criticor_output")
            
            critic_btn.click(fn=critic, inputs=[key, critic_model, critic_shots, critic_temperature, critic_number, critic_auit_directory], 
                                                outputs=critic_output, api_name="criticor")
            
            gr.Markdown("### Ranker")
            ranker_strategy = gr.Radio(label="Strategy", value="Default", choices=["Custom", "Default"])
            ranker_audit_location = gr.Textbox(label="Auditor Directory Location", value="", 
                                               info="Example = auditor_{audit_model}_{audit_temperature}_top{audit_top_k}_{audit_number}")
            ranker_critic_location = gr.Textbox(label="Critic Directory Location",value="", 
                                                info="Example = critic_{critic_model}_{critic_temperature}_{critic_number}_{critic_shots}")
            with gr.Row():
                with gr.Column(min_width=300, scale=2):
                    ranker_btn = gr.Button(value="Start Ranker", variant="primary", elem_id="ranker")
                with gr.Column(min_width=300, scale=2):
                    gr.Button(value="Start Ranker", variant="primary", visible=False)

            ranker_output = gr.Textbox(label="Ranker Status", value=ranker_output, interactive=False, visible=True, elem_id="ranker_output")
                
            ranker_btn.click(fn=ranker, inputs=[critic_model, critic_shots, critic_temperature, critic_number, critic_auit_directory, 
                                                ranker_strategy, ranker_audit_location, ranker_critic_location], 
                                                    outputs=ranker_output, api_name="ranker")
            

            gr.Button(value="You can checkout the results at path - ./src/logs/", link="https://huggingface.co/spaces/Asolanki97/GPTLens/tree/main/src/logs")
            key.change(fn=openAIKey, inputs=key, outputs=message)



ui.launch(share=True)