import gradio as gr
from gpt4all import GPT4All
from urllib.request import urlopen
import json
import time

url = "https://raw.githubusercontent.com/nomic-ai/gpt4all/main/gpt4all-chat/metadata/models3.json"
response = urlopen(url)
data_json = json.loads(response.read())

def model_choices():
    model_list = [data_json[i]['filename'] for i in range(len(data_json))]
    return model_list

model_description = {model['filename']: model['description'] for model in data_json}

def llm_intro(selected_model):
    return model_description.get(selected_model, "No description available for this model selection.")

def generate_text(input_text, selected_model):
    model = GPT4All(selected_model)
    output = model.generate(input_text, max_tokens=100)
    return output

# model_dropdown = gr.Dropdown(choices=model_choices(), 
#                              multiselect=False, 
#                              label="LLMs to choose from", 
#                              type="value", 
#                              value="orca-mini-3b-gguf2-q4_0.gguf")

# explanation = gr.Textbox(label="Model Description", interactive=True, lines=10)

# model_dropdown.change(fn=llm_intro, inputs=model_dropdown, outputs=explanation)

# input_text = gr.Textbox(lines=5, label="Input Text")
# output_text = gr.Textbox(lines=5, label="Generated Text")

# gr.Interface(fn=generate_text, 
#              inputs=[input_text, model_dropdown], 
#              outputs=output_text, 
#              theme = gr.themes.Soft(), 
#              analytics_enabled=True, 
#              title="GPT4All Text Generation Experiment").launch()

with gr.Blocks() as demo:
    gr.Markdown("## GPT4All Text Generation Experiment")

    with gr.Row():
        
    
        model_selection = gr.Dropdown(choices=model_choices(), 
                                     multiselect=False, 
                                     label="LLMs to choose from", 
                                     type="value", 
                                     value="orca-mini-3b-gguf2-q4_0.gguf")
    
        explanation = gr.Textbox(label="Model Description", lines=3, 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_selection.change(fn=llm_intro, inputs=model_selection, outputs=explanation)

    chatbot = gr.Chatbot()
    input_text = gr.Textbox(lines=3, label="Press shift+Enter to submit")
    # output_text = gr.Textbox(lines=10, label="Generated Text")
    clear = gr.ClearButton([input_text, chatbot])

    def respond(message, chat_history, selected_model):
        bot_message = generate_text(message, selected_model)
        chat_history.append((message, bot_message))
        time.sleep(2)  # Simulating processing delay if necessary
    
        return bot_message, chat_history

    input_text.submit(respond, [input_text, chatbot, model_selection], [chatbot])
        
    # # Button to generate text
    # generate_btn = gr.Button("Generate")
    # generate_btn.click(fn=generate_text, inputs=[input_text, model_selection], outputs=output_text)

demo.launch()