import gradio as gr
import os

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 = [
    'gpt2',
    'lm-human-preference-details/train_policy_accelerate_tf_adam_gpt2_xl_grad_accu__sentiment_offline_5k.json__seed2',
    'lm-human-preference-details/train_policy_accelerate_tf_adam_gpt2_xl_grad_accu__descriptiveness_offline_5k.json__seed1',
]


interfaces = [gr.load(f"models/{name}") for name in name_list]
def generate_code_os(text):
    return [interface(text) for interface in interfaces]


def set_example(example: list) -> dict:
    return gr.Textbox.update(value=example[0]) 
    
with gr.Blocks() as demo:
    gr.Markdown("Compare RLHF Stylistic Models")
    with gr.Box():
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(label = "Write your prompt here", lines=4)
                with gr.Row():
                    btn = gr.Button("Generate")
                
                # example_text = gr.Dataset(components=[input_text], samples=examples)
                # example_text.click(fn=set_example,
                #                 inputs = example_text,
                #                 outputs= example_text)    
                gr.Examples(
                    [["The young maid said"], ["The table has arrived at the house"]],
                    [input_text],
                )
            with gr.Column():   
                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)