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)