import gradio as gr
from gpt4all import GPT4All
from urllib.request import urlopen
import json
import time
from load_llms import model_choices, llm_intro, load_model

    
# Construct chatbot
def generate_response(model_name, message, chat_history):
    model = load_model(model_name)
    response = model.generate(message, max_tokens=100)
    chat_history.append((message, response))
    return "", chat_history

    
# Create Gradio UI
with gr.Blocks(
    css=".contain { display: flex !important; flex-direction: column !important; }"
    "#chatbot { flex-grow: 1 !important; overflow: auto !important; }"
    "#col { height: calc(100vh - 112px - 16px) !important; }"
    "#submit:hover { background-color: green !important; }"
    "#clear { background-color: darkred !important; }",
    theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Martel Sans")])                  
) as demo:
    gr.Markdown("# GPT4All Chatbot")
    with gr.Row():
        with gr.Column(scale=1):
            model_dropdown = gr.Dropdown(
                choices=model_choices(),  
                multiselect=False,
                type="value",
                value="orca-mini-3b-gguf2-q4_0.gguf",
                label="LLMs to choose from"
            )
            explanation = gr.Textbox(label="Model Description", interactive=False, value=llm_intro("orca-mini-3b-gguf2-q4_0.gguf"))
    
            # Link the dropdown with the textbox to update the description based on the selected model
            model_dropdown.change(fn=llm_intro, inputs=model_dropdown, outputs=explanation)

            gr.Image(interactive=False, value="https://huggingface.co./spaces/yxmauw/chatbot/blob/main/space_hardware.png", type="filepath", image_mode="RGBA")
            
        with gr.Column(scale=4, elem_id='col'):
            chatbot = gr.Chatbot(label="Chatroom", value=[(None, "How may I help you today?")], elem_id="chatbot")

            with gr.Row():       
                message = gr.Textbox(label="Message", scale=10)
                message.submit(generate_response, inputs=[model_dropdown, message, chatbot], outputs=[message, chatbot])
                submit_button = gr.Button("Submit", scale=1, elem_id="submit")
                submit_button.click(generate_response, inputs=[model_dropdown, message, chatbot], outputs=[message, chatbot])

            clear = gr.ClearButton([message, chatbot], elem_id="clear")

# Launch the Gradio app
demo.launch()