import gradio as gr
from openai import OpenAIError

from selfinstruct.instance import instance_main
import openai
from langchain.llms.openai import OpenAI
import os
import json

def OpenaiChat(key, input, temperature):
    print(input)
    openai.api_key = key
    MODEL = "gpt-3.5-turbo"
    try:
        response = openai.ChatCompletion.create(
            model=MODEL,
            messages=[
                {"role": "user", "content": input}
            ],
            temperature=temperature
        )
        return json.dumps(response["choices"], ensure_ascii=False)
    except OpenAIError:
        return "Wrong Key!"

with gr.Blocks() as demo:
    gr.Markdown("# Self-instruct Prompt")
    with gr.Row():
        key = gr.Textbox(label="Input OpenAI key", type="password")
    with gr.Tab("Generate instance"):
        with gr.Row():
            with gr.Column():
                gr.Dropdown(["davinci"], value=["davinci"], multiselect=False, label="Model")
                text_input = gr.TextArea(lines=3, label="Input instruction:", info="格式：一行一个instruction")
            with gr.Column():
                text_output = gr.JSON(lines=10, label="Output instance:")
        text_button = gr.Button("Generate")
        gr.Markdown("## 3 Steps to generate")
        with gr.Row():
            text_step1 = gr.TextArea(lines=3, label="Step1 identify if classifacrion task:", )
            text_step2 = gr.TextArea(lines=3, label="Step2 generate instance:", )
            text_step3 = gr.TextArea(lines=3, label="Step3 Prepare for finetune:", )

    with gr.Tab("Simple ChatGPT"):
        gr.Dropdown(["GPT3.5"], value=["GPT3.5"], multiselect=False, label="Model")
        temperature = gr.Slider(0, 1.0, value=0.5, label="Temperature", info="Choose betwen 0 and 1")
        ChatGPT_input = gr.TextArea(lines=3, label="Input:")
        ChatGPT_output = gr.TextArea(lines=3, label="Output:")
        ChatGPT_button = gr.Button("Run")

    # with gr.Accordion("Open for More!"):
    #     gr.Markdown("Look at me...")

    text_button.click(instance_main, inputs=[text_input, key], outputs=[text_output, text_step1, text_step2, text_step3])
    ChatGPT_button.click(OpenaiChat, inputs=[key, ChatGPT_input, temperature], outputs=ChatGPT_output)

demo.launch()