mushroomsolutions/TestDataGeneration

1import gradio as gr2import pandas as pd 3from realtabformer import REaLTabFormer4from scipy.io import arff5import os6 7rtf_model = REaLTabFormer(8    model_type="tabular",9    epochs=25, # Default is 20010    gradient_accumulation_steps=4)11 12 13def generate_data(file, num_samples):14    if '.arff' in file.name:15        data = arff.loadarff(open(file.name,'rt'))16        df = pd.DataFrame(data[0])17    elif '.csv' in file.name:18        df = pd.read_csv(file.name)19    rtf_model.fit(df, num_bootstrap=10) # Default is 50020    # Generate synthetic data21    samples = rtf_model.sample(n_samples=num_samples)22 23    return samples24 25def generate_relational_data(parent_file, child_file, join_on):26    parent_df = pd.read_csv(parent_file.name)27    child_df = pd.read_csv(child_file.name)28 29    #Make sure join_on column exists in both30    assert ((join_on in parent_df.columns) and31        (join_on in child_df.columns))32 33    rtf_model.fit(parent_df.drop(join_on, axis=1), num_bootstrap=100)34 35    pdir = Path("rtf_parent/")36    rtf_model.save(pdir)37 38    # # Get the most recently saved parent model,39    # # or a specify some other saved model.40    # parent_model_path = pdir / "idXXX"41    parent_model_path = sorted([42        p for p in pdir.glob("id*") if p.is_dir()],43        key=os.path.getmtime)[-1]44 45    child_model = REaLTabFormer(46    model_type="relational",47    parent_realtabformer_path=parent_model_path,48    epochs = 25,49    output_max_length=None,50    train_size=0.8)51 52    child_model.fit(53    df=child_df,54    in_df=parent_df,55    join_on=join_on,56    num_bootstrap=10)57 58    # Generate parent samples.59    parent_samples = rtf_model.sample(5)60 61    # Create the unique ids based on the index.62    parent_samples.index.name = join_on63    parent_samples = parent_samples.reset_index()64 65    # Generate the relational observations.66    child_samples = child_model.sample(67        input_unique_ids=parent_samples[join_on],68        input_df=parent_samples.drop(join_on, axis=1),69        gen_batch=5)70 71    return parent_samples, child_samples, gr.update(visible = True)72    73 74with gr.Blocks() as demo:75    gr.Markdown("""76                ## REaLTabFormer: Generating Realistic Relational and Tabular Data using Transformers77            """)78    gr.HTML('''79     <p style="margin-bottom: 10px; font-size: 94%">80                This is an unofficial demo for REaLTabFormer, an approach that can be used to generate synthetic data from single tabular data using GPT. The demo is based on the <a href='https://github.com/avsolatorio/REaLTabFormer' style='text-decoration: underline;' target='_blank'> Github </a> implementation provided by the authors.81              </p>82              ''')83    gr.HTML('''84    <p align="center"><img src="https://github.com/avsolatorio/RealTabFormer/raw/main/img/REalTabFormer_Final_EQ.png" style="width:40%"/></p>85    ''')86    87    with gr.Column():88        89        with gr.Tab("Upload Data as File: Tabular Data"):90            data_input_u = gr.File(label = 'Upload Data File (Currently supports CSV and ARFF)', file_types=[".csv", ".arff"])91            num_samples = gr.Slider(label="Number of Samples", minimum=5, maximum=100, value=5, step=10)92            generate_data_btn = gr.Button('Generate Synthetic Data')93 94        with gr.Tab("Upload Data as File: Relational Data"):95            data_input_parent = gr.File(label = 'Upload Data File for Parent Dataset', file_types=[ ".csv"])96            data_input_child = gr.File(label = 'Upload Data File for Child Dataset', file_types=[ ".csv"])97            join_on = gr.Textbox(label = 'Column name to join on')98            99            generate_data_btn_relational = gr.Button('Generate Synthetic Data')100 101        with gr.Row():102            #data_sample = gr.Dataframe(label = "Original Data")103            data_output = gr.Dataframe(label = "Synthetic Data")104        with gr.Row(visible = False) as child_sample:105            data_output_child = gr.Dataframe(label = "Synthetic Data for Child Dataset")106    107    108    generate_data_btn.click(generate_data, inputs = [data_input_u,num_samples], outputs = [data_output])109    generate_data_btn_relational.click(generate_relational_data, inputs = [data_input_parent,data_input_child,join_on], outputs = [data_output, data_output_child, child_sample])110    examples = gr.Examples(examples=[['diabetes.arff',5], ["titanic.csv", 15]],inputs = [data_input_u,num_samples], outputs = [data_output], cache_examples = True, fn = generate_data)111 112    113demo.launch()