import gradio as gr
import os
import openai

openai.api_key = "API"

demo = gr.Blocks()

#'spaces/Gradio-Blocks/Codex_OpenAI' did not work due to paramerts error 
#'huggingface/facebook/incoder-6B' inference 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']

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="",
    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="",
    temperature=0,
    max_tokens=256,
    top_p=1,
    frequency_penalty=0,
    presence_penalty=0
    ).choices[0].text
    return response_2


with gr.Blocks() as demo:
    gr.Markdown("<h1><center>Compare OpenAI and HuggingFace Code Generators</center></h1>")
    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 ...")
                    
            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(debug=True)