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 = "Nearly a half-century ago, Apollo 11 astronaut Neil Armstrong walked on the Moon. NASA is now preparing for an ambitious new era of sustainable human spaceflight and discovery. The agency is building the Space Launch System rocket and the Orion spacecraft for human deep space exploration. With the help of commercial and international partners, NASA will develop new opportunities in lunar orbit, including a platform to aid surface exploration and serve as a gateway to Mars." # 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()