import gradio as gr

# Import necessary modules
from tokenizer.tokenizer import *
from models.GPT import build_GPT
from inference.inference import Generative_inference
from inference.sampling_strategies.sample_random import *

# Initialize tokenizer
tokenizer = SPM_Tokenizer(vocab_model_file='./tokenizer_.model', input_size=256+1)

# Define model parameters
vocab_size = 454+1
input_len = 256

# Build GPT model
GPT, flops = build_GPT(256, vocab_size, 1000, 2, 0, 50, 20, 5)

# Load pre-trained weights
GPT.load_weights('AEON_30M.weights.h5')

# Inference stance
inference = Generative_inference(GPT, tokenizer, input_len=256, k_value=5)

# Default text
default_input_text = "A black hole is a region of spacetime where gravity is so strong that nothing, including light and other electromagnetic waves, is capable of possessing enough energy to escape it. Einstein's theory of general relativity predicts that a sufficiently compact mass can deform spacetime to form a black hole. The boundary of no escape is called the event horizon. A black hole has a great effect on the fate and circumstances of an object crossing it, but it has no locally detectable features according to general relativity. In many ways, a black hole acts like an ideal black body, as it reflects no light."

# Define the function to generate text based on input
def generate_text(input_text, k_value=10, generate_limit=50):
    generated_text = inference.generate(input_text, generate_limit=generate_limit, k_value=k_value)
    return generated_text

# Create Gradio interface blocks
with gr.Blocks() as demo:
    # Warning message for users
    gr.Markdown("### Warning")
    gr.Markdown("The current model is not a conversational or text domain-specific model. It was trained on a range of essays and articles on space, religion, philosophy, and current affairs. It works as a text autocompleting model that can be used to fine-tune for different purposes. Enter a text with 100 words or copy-paste it here for the best results.")

    # Model specifications
    gr.Markdown("### Model Specs")
    gr.Markdown("This is a 30M parameter model ONLY! This tiny model is free to use for any purpose under the Apache 2.0 license. Once quantized, it can work on mobile CPUs too for tiny language model purposes.")

    # Input components: text input, sliders for k_value and generate_limit
    with gr.Row():
        input_text = gr.Textbox(label="Input Text", lines=10, value=default_input_text)
        k_value = gr.Slider(minimum=1, maximum=30, value=10, step=1, label="K Value")
        generate_limit = gr.Slider(minimum=1, maximum=500, value=50, step=1, label="Generate Limit")
    
    # Output component: text output
    output_text = gr.Textbox(label="Output Text")

    # Button to trigger text generation
    generate_button = gr.Button("Generate")
    generate_button.click(generate_text, inputs=[input_text, k_value, generate_limit], outputs=output_text)

# Launch the Gradio interface
demo.launch()