import gradio as gr
import os
import openai

openai.api_key = os.environ["API"]

demo = gr.Blocks()

#'spaces/Gradio-Blocks/Codex_OpenAI' did not work due to paramerts error 
#'huggingface/facebook/incoder-6B' inference error
#huggingface/Salesforce/codegen-2B-mono 500 Error 

name_list = ['huggingface/lvwerra/codeparrot','huggingface/facebook/incoder-1B','huggingface/Salesforce/codegen-2B-mono']

open_ai_name = ['openai/code-davinci-002', 'openai/code-cushman-001']

examples = [ ["import numpy as np\nfrom sklearn.ensemble import RandomForestClassifier\n# create training data\nX = np.random.randn(100, 100)\ny = np.random.randint(0, 1, 100)\n # setup train test split"], 
             ['def get_file_size(filepath):'],
            ] 

def generate_code_os(text):
    interfaces = [gr.Interface.load(name) for name in name_list]
    return [interface(text) for interface in interfaces]
     
def generate_code_cs_1(text):
    response = openai.Completion.create(
    engine="code-davinci-002",
    prompt=text,
    temperature=0,
    max_tokens=256,
    top_p=1,
    frequency_penalty=0,
    presence_penalty=0
    ).choices[0].text
    response_2 = generate_code_cs_2(text)
    return [response, response_2]

def generate_code_cs_2(text):
    response_2 = openai.Completion.create(
    engine="code-cushman-001",
    prompt=text,
    temperature=0,
    max_tokens=256,
    top_p=1,
    frequency_penalty=0,
    presence_penalty=0
    ).choices[0].text
    return response_2

def set_example(example: list) -> dict:
    return gr.Textbox.update(value=example[0]) 
    
with gr.Blocks() as demo:
    gr.Markdown("# Compare OpenAI and HuggingFace Code Generators")
    with gr.Box():
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(label = "Write your code here", lines=4)
                with gr.Row():
                    btn = gr.Button("Open-Source models Generate code brrr ...")
                    btn2 = gr.Button("OpenAI models Generate code brrr ...")
                
                example_text = gr.Dataset(components=[input_text], samples=examples)
                example_text.click(fn=set_example,
                                inputs = example_text,
                                outputs= example_text.components)    
            with gr.Column():   
                gr.Markdown("Let’s Compare code generators ") 
                btn2.click(generate_code_cs_1, inputs = input_text, outputs =  [gr.Textbox(label=open_ai_name[_], lines=4) for _ in range(len(open_ai_name))])
                with gr.Column():       
                    btn.click(generate_code_os, inputs = input_text, outputs = [gr.Textbox(label=name_list[_], lines=4) for _ in range(len(name_list))])


demo.launch(enable_queue=True,debug=True)