Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import pandas as pd 
 from realtabformer import REaLTabFormer
 from scipy.io import arff
+import os
 
 rtf_model = REaLTabFormer(
     model_type="tabular",
@@ -20,7 +21,57 @@ def generate_data(file, num_samples):
     samples = rtf_model.sample(n_samples=num_samples)
 
     return samples
+
+def generate_relational_data(parent_file, child_file, join_on):
+    parent_df = pd.read_csv(parent_file.name)
+    child_df = pd.read_csv(child_file.name)
+
+    #Make sure join_on column exists in both
+    assert ((join_on in parent_df.columns) and
+        (join_on in child_df.columns))
+
+    rtf_model.fit(parent_df.drop(join_on, axis=1), num_bootstrap=100)
+
+    pdir = Path("rtf_parent/")
+    rtf_model.save(pdir)
+
+    # # Get the most recently saved parent model,
+    # # or a specify some other saved model.
+    # parent_model_path = pdir / "idXXX"
+    parent_model_path = sorted([
+        p for p in pdir.glob("id*") if p.is_dir()],
+        key=os.path.getmtime)[-1]
+
+    child_model = REaLTabFormer(
+    model_type="relational",
+    parent_realtabformer_path=parent_model_path,
+    epochs = 25,
+    output_max_length=None,
+    train_size=0.8)
+
+    child_model.fit(
+    df=child_df,
+    in_df=parent_df,
+    join_on=join_on,
+    num_bootstrap=100)
+
+    # Generate parent samples.
+    parent_samples = rtf_model.sample(5)
+
+    # Create the unique ids based on the index.
+    parent_samples.index.name = join_on
+    parent_samples = parent_samples.reset_index()
+
+    # Generate the relational observations.
+    child_samples = child_model.sample(
+        input_unique_ids=parent_samples[join_on],
+        input_df=parent_samples.drop(join_on, axis=1),
+        gen_batch=5)
+
+    return parent_samples, child_samples, gr.update(visible = True)
     
+
+
 css = """
         .gradio-container {
             font-family: 'IBM Plex Sans', sans-serif;
@@ -106,7 +157,7 @@ with gr.Blocks(css = css) as demo:
             """)
     gr.HTML('''
      <p style="margin-bottom: 10px; font-size: 94%">
-                This is an unofficial demo for REaLTabFormer 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.
+                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.
               </p>
               ''')
     
@@ -116,18 +167,27 @@ with gr.Blocks(css = css) as demo:
         #     audio_input_r = gr.Audio(label = 'Record Audio Input',source="microphone",type="filepath")
         #     transcribe_audio_r = gr.Button('Transcribe')
         
-        with gr.Tab("Upload Data as File"):
-            data_input_u = gr.File(label = 'Upload Data File', file_types=["text", ".json", ".csv", ".arff"])
+        with gr.Tab("Upload Data as File: Tabular Data"):
+            data_input_u = gr.File(label = 'Upload Data File (Currently supports CSV and ARFF)', file_types=[".csv", ".arff"])
             num_samples = gr.Slider(label="Number of Samples", minimum=5, maximum=100, value=5, step=10)
             generate_data_btn = gr.Button('Generate Synthetic Data')
 
+        with gr.Tab("Upload Data as File: Relational Data"):
+            data_input_parent = gr.File(label = 'Upload Data File for Parent Dataset', file_types=[ ".csv"])
+            data_input_child = gr.File(label = 'Upload Data File for Child Dataset', file_types=[ ".csv"])
+            join_on = gr.Textbox(label = 'Column name to join on')
+            
+            generate_data_btn_relational = gr.Button('Generate Synthetic Data')
+
         with gr.Row():
             #data_sample = gr.Dataframe(label = "Original Data")
             data_output = gr.Dataframe(label = "Synthetic Data")
-            
+        with gr.Row(visible = False) as child_sample:
+            data_output_child = gr.Dataframe(label = "Synthetic Data for Child Dataset")
     
     
     generate_data_btn.click(generate_data, inputs = [data_input_u,num_samples], outputs = [data_output])
+    generate_data_btn_relational.click(generate_relational_data, inputs = [data_input_parent,data_input_child,join_on], outputs = [data_output, data_output_child])
     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)