Update app.py

app.py

@@ -1,5 +1,6 @@
 import io
 import random
+import tempfile
 from dataclasses import dataclass
 
 import gradio as gr
@@ -64,6 +65,29 @@ def fig_to_image(fig) -> np.ndarray:
     return image
 
 
+def build_full_dataset_df(X_train, y_train, X_val, y_val) -> pd.DataFrame:
+    """Create a single DataFrame with a 'split' column so it’s easy to teach/train/export."""
+    cols = [f"x{i}" for i in range(10)]
+
+    train_df = pd.DataFrame(X_train.cpu().numpy(), columns=cols)
+    train_df["y"] = y_train.cpu().numpy().reshape(-1)
+    train_df["split"] = "train"
+
+    val_df = pd.DataFrame(X_val.cpu().numpy(), columns=cols)
+    val_df["y"] = y_val.cpu().numpy().reshape(-1)
+    val_df["split"] = "val"
+
+    full_df = pd.concat([train_df, val_df], axis=0, ignore_index=True)
+    return full_df
+
+
+def save_df_to_temp_csv(df: pd.DataFrame) -> str:
+    """Save DataFrame to a temp CSV and return the file path for Gradio download."""
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="synthetic_linear_regression_")
+    df.to_csv(tmp.name, index=False)
+    return tmp.name
+
+
 def train_raw_pytorch(
     n_samples: int,
     noise_std: float,
@@ -79,6 +103,10 @@ def train_raw_pytorch(
     spec = DataSpec(n_samples=n_samples, n_features=10, noise_std=noise_std, train_frac=0.8)
     X_train, y_train, X_val, y_val, w_true, b_true = make_synthetic_regression(spec, seed=seed)
 
+    # Full dataset CSV (train + val with split column)
+    full_df = build_full_dataset_df(X_train, y_train, X_val, y_val).round(4)
+    csv_path = save_df_to_temp_csv(full_df)
+
     # Data preview (first 20 rows from training split)
     preview_n = min(20, X_train.shape[0])
     df_preview = pd.DataFrame(
@@ -205,7 +233,8 @@ def train_raw_pytorch(
         f"Final train loss: {train_losses[-1]:.6f}\n"
         f"Final val loss:   {val_losses[-1]:.6f}\n\n"
         f"True bias (b_true):        {float(b_true.item()):.4f}\n"
-        f"Learned bias (b_learned):  {float(b_learned.item()):.4f}\n"
+        f"Learned bias (b_learned):  {float(b_learned.item()):.4f}\n\n"
+        f"Dataset CSV includes columns: x0..x9, y, split(train/val)\n"
     )
 
     raw_loop_snippet = """# Raw PyTorch: requires manual training loop
@@ -224,7 +253,8 @@ for x, y in dataloader:
     optimizer.step()
 """
 
-    return loss_plot, df_weights, summary, raw_loop_snippet, df_preview
+    # Added csv_path as downloadable artifact
+    return loss_plot, df_weights, summary, raw_loop_snippet, df_preview, csv_path
 
 
 with gr.Blocks(title="Raw PyTorch Training Loop Demo") as demo:
@@ -238,7 +268,7 @@ This Space generates **synthetic data** each run:
 y = Xw + b + \\text{noise}
 \]
 
-Use **Data Preview** to see the first rows of the generated dataset.
+Go to **Data Preview** to see sample rows and **download the full dataset** as CSV.
 """
     )
 
@@ -261,10 +291,11 @@ Use **Data Preview** to see the first rows of the generated dataset.
     with gr.Tab("Outputs"):
         loss_img = gr.Image(label="Loss Curve", type="numpy")
         weights_df = gr.Dataframe(label="Weights: True vs Learned (sorted by abs error)", wrap=True)
-        summary_txt = gr.Textbox(label="Summary", lines=8)
+        summary_txt = gr.Textbox(label="Summary", lines=10)
 
     with gr.Tab("Data Preview"):
-        data_preview = gr.Dataframe(label="First 20 rows of the generated training data (X features + y)", wrap=True)
+        data_preview = gr.Dataframe(label="First 20 rows of generated TRAIN data (X features + y)", wrap=True)
+        download_file = gr.File(label="Download full dataset CSV (train + val)")
 
     with gr.Tab("Raw Loop Snippet"):
         snippet = gr.Code(label="Your original loop (as runnable reference)", language="python")
@@ -272,7 +303,7 @@ Use **Data Preview** to see the first rows of the generated dataset.
     run_btn.click(
         fn=train_raw_pytorch,
         inputs=[n_samples, noise_std, lr, batch_size, epochs, seed, device_choice],
-        outputs=[loss_img, weights_df, summary_txt, snippet, data_preview],
+        outputs=[loss_img, weights_df, summary_txt, snippet, data_preview, download_file],
     )
 
 if __name__ == "__main__":