import gradio as gr

from keras_hub.models import GemmaCausalLM

# # Path to your locally saved model and config
# model_weights_path = "gemma_finetuned_model/finetunedmodel.weights.h5"
# config_path = "gemma_finetuned_model/config.json"
#
# # Load the model using the local config and weights
# with open(config_path, "r") as f:
#     config = json.load(f)
#
# gemma_lm = GemmaCausalLM.from_config(config)
# gemma_lm.load_weights(model_weights_path)
gemma_lm = GemmaCausalLM.from_preset("hf://bhaveshgoel07/MedCode")


#
# # Load the fine-tuned model
# gemma_lm = GemmaCausalLM.from_preset("gemma_2b_en", **config)
# gemma_lm.load_weights("gemma_finetuned_model/model_weights.h5")

# Gradio app function
def generate_response(prompt):
    template = "Instruction:\n{instruction}\n\nResponse:\n{response}"
    formatted_prompt = template.format(instruction=prompt, response="")

    # Generate output
    output = gemma_lm.generate(formatted_prompt, max_length=100)

    return output


# Use Gradio Blocks for more layout control
with gr.Blocks() as interface:
    gr.Markdown("## Fine-tuned GemmaCausalLM Chatbot")
    gr.Markdown("Enter a prompt to get a response from the fine-tuned GemmaCausalLM model.")

    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt")
        response_output = gr.Textbox(label="Response")

    generate_button = gr.Button("Generate")
    generate_button.click(generate_response, inputs=prompt_input, outputs=response_output)

    # Add footer text at the bottom
    gr.Markdown("Some Examples:\n 1. Cholera due to Vibrio cholerae 01, biovar cholerae ICD Code:A00.0\n 2. Cholera, Unspecified ICD Code:A00.9\n 3. Typhoid pneumonia, Typhoid Fever ICD Code:A01.03")

# Launch the interface
interface.launch()