app.py

import gradio as gr
from defaults import DEFAULTS

def greet(name, intensity)->str:
    return "Hello, " + name + "!" * int(intensity)

def create_parallelism_block():
    with gr.Column():
        gr.Markdown("# Parallelism Parameters")
        tp = gr.Number(label="Tensor Parallelism", value=1, interactive=True)
        pp = gr.Number(label="Pipeline Parallelism", value=1, interactive=True)
        cp = gr.Number(label="Context Parallelism", value=1, interactive=True)
        ep = gr.Number(label="Expert Parallelism", value=1, interactive=True)
        return tp, pp, cp, ep

def create_model_block():
    with gr.Column():
        gr.Markdown("# Model Parameters")
        layers = gr.Number(label="Number of Layers", value=32, interactive=True)
        vocab = gr.Number(label="Vocab Size", value=32000, interactive=True)
        hidden = gr.Number(label="Hidden Dim", value=4096, interactive=True)
        intermediate = gr.Number(label="Intermediate Dim", value=11008, interactive=True)
        presets = gr.Dropdown(list(DEFAULTS.keys()), label="Presets", interactive=True)
        return layers, vocab, hidden, intermediate, presets

def create_training_block():
    with gr.Column():
        gr.Markdown('# Training Parameters')
        seq_len = gr.Number(label="Sequence Length", value=8192, interactive=True)
        batch_size = gr.Number(label="Batch Size", value=8, interactive=True)
        return seq_len, batch_size

def calculate(*args)->int:
    out = 1
    for arg in args:
        out *= arg
    return arg
    


with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            tp, pp, cp, ep = create_parallelism_block()
            layers, vocab, hidden, intermediate, presets = create_model_block()
            seq_len, batch_size = create_training_block()
        calculate_button = gr.Button("Calculate")
        output = gr.Number(label="Output")
        
        calculate_button.click(fn=calculate, inputs=[tp,pp,cp,ep],outputs=output)
        

demo.launch()