Update app.py

app.py

@@ -1,826 +1,346 @@
+# --- Standard Library Imports ---
+import os
+import time
+import traceback
+import tempfile
+import json
+import math
+import collections
+import collections.abc # For Gradio compatibility with newer Python versions
+
+# --- UI Framework ---
 import gradio as gr
+
+# --- Data Handling & Numerical Ops ---
 import pandas as pd
 import numpy as np
-import sklearn
+
+# --- Core Machine Learning (Scikit-learn) ---
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import StandardScaler, LabelEncoder, OneHotEncoder
 from sklearn.impute import SimpleImputer
 from sklearn.compose import ColumnTransformer
 from sklearn.pipeline import Pipeline
-# Scikit-learn Models
 from sklearn.linear_model import LogisticRegression, LinearRegression
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
 from sklearn.svm import SVC, SVR
-# Metrics
 from sklearn.metrics import accuracy_score, classification_report, mean_squared_error, r2_score
-# Dataset generators
 from sklearn.datasets import make_classification, make_regression
-
 import joblib
-import os
-import time
+
+# --- Core Machine Learning (PyTorch) ---
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import TensorDataset, DataLoader
-import torchvision # For transforms, even if data is basic
-import torchvision.transforms as T
 
-# ONNX specific imports
+# --- ONNX Support for Model Interoperability ---
 import skl2onnx
 from skl2onnx import convert_sklearn
-from skl2onnx.common.data_types import FloatTensorType, Int64TensorType, StringTensorType # Ensure this import if used for ONNX initial types
-import onnxruntime as rt
+from skl2onnx.common.data_types import FloatTensorType, StringTensorType
 
-import traceback
-import tempfile
-import json
-import math
-import collections.abc # For Gradio issue with new Python versions
-import collections # Added for OrderedDict if not already covered
-import matplotlib # Use Agg backend for non-interactive environments
-matplotlib.use('Agg')
+# --- Visualization ---
+import matplotlib
+matplotlib.use('Agg') # Use non-interactive backend for server environments
 import matplotlib.pyplot as plt
 
+# --- Graceful ONNX Runtime Handling ---
+# This addresses the system-level ImportError on platforms like Hugging Face Spaces.
+try:
+    import onnxruntime as rt
+    ONNX_RUNTIME_AVAILABLE = True
+except ImportError:
+    ONNX_RUNTIME_AVAILABLE = False
+    print("Warning: onnxruntime could not be imported. ONNX model validation will be skipped.")
+# --- End of Imports ---
+
 
-# --- Global Variables / Constants ---
+# --- Global Variables & Constants ---
 TEMP_DIR = "temp_outputs"
 os.makedirs(TEMP_DIR, exist_ok=True)
-MAX_DATASET_ROWS_WARN = 30000
 MAX_GENERATED_ROWS = 50000 
 MAX_GENERATED_COLS = 100   
+PARAM_RANGES = collections.OrderedDict([
+    ("Tiny (<10k)", (0, 10000)),
+    ("Small (10k-50k)", (10000, 50000)),
+    ("Medium (50k-250k)", (50000, 250000)),
+    ("Large (250k-1M)", (250000, 1000000)),
+])
 
-# --- Helper Functions (count_parameters, get_temp_filepath) ---
-def count_sklearn_parameters(model):
-    if hasattr(model, 'coef_'):
-        return model.coef_.size + (model.intercept_.size if hasattr(model, 'intercept_') else 0)
-    if hasattr(model, 'support_vectors_'):
-        return model.support_vectors_.size
-    if isinstance(model, (RandomForestClassifier, RandomForestRegressor)):
-        try:
-            return sum(tree.tree_.node_count for tree in model.estimators_)
-        except: return "N/A (Complex Ensemble)"
-    return "N/A"
-
-def count_pytorch_parameters(model):
-    return sum(p.numel() for p in model.parameters() if p.requires_grad)
-
+# --- Helper Functions ---
 def get_temp_filepath(filename_base, extension):
+    """Generates a unique temporary filepath."""
     clean_extension = extension.lstrip('.')
     return os.path.join(TEMP_DIR, f"{filename_base}_{time.strftime('%Y%m%d-%H%M%S')}.{clean_extension}")
 
-# --- PyTorch Model Definitions (SimpleMLP, SimpleCNN) ---
+# --- PyTorch Model Definitions ---
 class SimpleMLP(nn.Module):
+    """A simple Multi-Layer Perceptron."""
     def __init__(self, input_dim, hidden_layers_str, output_dim, activation_fn_str="relu", task_type="classification"):
-        super(SimpleMLP, self).__init__()
+        super().__init__()
         layers = []
-        if not isinstance(input_dim, int) or input_dim <= 0:
-            raise ValueError(f"Input dimension must be a positive integer, got {input_dim}")
-
-        hidden_units_list = []
-        if hidden_layers_str and isinstance(hidden_layers_str, str) and hidden_layers_str.strip():
-            try:
-                hidden_units_list = [int(x.strip()) for x in hidden_layers_str.split(',') if x.strip()]
-                if any(h_units <= 0 for h_units in hidden_units_list):
-                    raise ValueError("Hidden layer units must be positive integers.")
-            except ValueError as e:
-                raise ValueError(f"Invalid hidden layer string '{hidden_layers_str}'. Error: {e}")
+        hidden_units = [int(x.strip()) for x in hidden_layers_str.split(',') if x.strip()]
         
         current_dim = input_dim
-        for h_units in hidden_units_list:
+        for h_units in hidden_units:
             layers.append(nn.Linear(current_dim, h_units))
             if activation_fn_str.lower() == "relu": layers.append(nn.ReLU())
             elif activation_fn_str.lower() == "tanh": layers.append(nn.Tanh())
             elif activation_fn_str.lower() == "sigmoid": layers.append(nn.Sigmoid())
-            else: layers.append(nn.ReLU())
             current_dim = h_units
         
         layers.append(nn.Linear(current_dim, output_dim))
         
-        if task_type == "classification":
-            if output_dim == 1: # For BCELoss (binary classification)
-                layers.append(nn.Sigmoid())
-            # For multi-class, nn.CrossEntropyLoss expects raw logits, so no final activation here.
-        self.network = nn.Sequential(*layers)
-    def forward(self, x): return self.network(x)
-
-class SimpleCNN(nn.Module):
-    def __init__(self, input_channels, img_size_wh, num_classes=10, task_type="classification",
-                 c_out1=16, k1=3, s1=1, p1=1, pool1_k=2, pool1_s=2,
-                 c_out2=32, k2=3, s2=1, p2=1, pool2_k=2, pool2_s=2,
-                 fc_hidden=128):
-        super(SimpleCNN, self).__init__()
-        self.input_channels = input_channels
-        self.img_h, self.img_w = img_size_wh
-        self.num_classes = num_classes
-
-        self.conv1 = nn.Conv2d(self.input_channels, c_out1, kernel_size=k1, stride=s1, padding=p1)
-        self.relu1 = nn.ReLU()
-        self.pool1 = nn.MaxPool2d(kernel_size=pool1_k, stride=pool1_s)
-
-        h_out_conv1 = (self.img_h - k1 + 2 * p1) // s1 + 1
-        w_out_conv1 = (self.img_w - k1 + 2 * p1) // s1 + 1
-        h_pool1 = (h_out_conv1 - pool1_k) // pool1_s + 1
-        w_pool1 = (w_out_conv1 - pool1_k) // pool1_s + 1
-
-        self.conv2 = nn.Conv2d(c_out1, c_out2, kernel_size=k2, stride=s2, padding=p2)
-        self.relu2 = nn.ReLU()
-        self.pool2 = nn.MaxPool2d(kernel_size=pool2_k, stride=pool2_s)
-
-        h_out_conv2 = (h_pool1 - k2 + 2 * p2) // s2 + 1
-        w_out_conv2 = (w_pool1 - k2 + 2 * p2) // s2 + 1
-        h_pool2 = (h_out_conv2 - pool2_k) // pool2_s + 1
-        w_pool2 = (w_out_conv2 - pool2_k) // pool2_s + 1
-        
-        self.flattened_size = c_out2 * h_pool2 * w_pool2
-        if self.flattened_size <= 0:
-            raise ValueError(f"Calculated flattened size is {self.flattened_size}. Check CNN params and image size.")
-        
-        self.fc1 = nn.Linear(self.flattened_size, fc_hidden)
-        self.relu3 = nn.ReLU()
-        self.fc2 = nn.Linear(fc_hidden, self.num_classes)
+        if task_type == "classification" and output_dim == 1:
+            layers.append(nn.Sigmoid()) # For BCELoss
+        # For multi-class, CrossEntropyLoss expects raw logits, so no final activation.
         
-        # The final activation is now a separate attribute for clarity.
-        if task_type == "classification":
-            if self.num_classes == 1: # Binary classification with BCELoss
-                self.final_activation = nn.Sigmoid()
-            else: # Multi-class classification with CrossEntropyLoss
-                self.final_activation = nn.Identity() # The loss function combines Softmax and NLLLoss.
-        else: # Regression
-            self.final_activation = nn.Identity()
+        self.network = nn.Sequential(*layers)
 
     def forward(self, x):
-        x = self.pool1(self.relu1(self.conv1(x)))
-        x = self.pool2(self.relu2(self.conv2(x)))
-        x = x.view(-1, self.flattened_size)
-        x = self.relu3(self.fc1(x))
-        x = self.fc2(x)
-        x = self.final_activation(x)
-        return x
-
-# --- Parameter Target Helpers ---
-PARAM_RANGES = collections.OrderedDict([
-    ("Tiny (<10k)", (0, 10000)),
-    ("Small (10k-50k)", (10000, 50000)),
-    ("Medium (50k-250k)", (50000, 250000)),
-    ("Large (250k-1M)", (250000, 1000000)),
-])
-
-def suggest_mlp_layers_for_range(input_dim, output_dim, target_range_str, current_logs=""):
-    logs = current_logs
-    if not target_range_str or target_range_str not in PARAM_RANGES:
-        logs += "Invalid parameter range selected for MLP suggestion.\n"; return "", logs
-    min_p, max_p = PARAM_RANGES[target_range_str]
-    target_p_avg = (min_p + max_p) // 2
-    suggested_layers_str = ""
-    if input_dim <=0 or output_dim <=0:
-        logs += "Input/Output dims must be positive for MLP suggestion.\n"; return "", logs
-
-    h1_candidate = max(1, int(target_p_avg / (input_dim + output_dim + 1e-6)))
-    params_1_layer = (input_dim * h1_candidate + h1_candidate) + (h1_candidate * output_dim + output_dim)
-    if min_p <= params_1_layer <= max_p and h1_candidate > 0:
-        suggested_layers_str = str(h1_candidate)
-        logs += f"Suggested 1 hidden layer: {h1_candidate} units (Est. Params: {params_1_layer})\n"
-    else:
-        h_base = max(1, int(math.sqrt(target_p_avg / 2.0)))
-        h1 = min(2048, max(1, int(h_base * (input_dim / (input_dim + output_dim + 1e-6)) * 2 + h_base / 2)))
-        h2 = min(2048, max(1, int(h_base * (output_dim / (input_dim + output_dim + 1e-6)) * 2 + h_base / 2)))
-        params_2_layers = (input_dim * h1 + h1) + (h1 * h2 + h2) + (h2 * output_dim + output_dim)
-        if min_p <= params_2_layers <= max_p and h1 > 0 and h2 > 0:
-            suggested_layers_str = f"{h1},{h2}"
-            logs += f"Suggested 2 hidden layers: {h1},{h2} units (Est. Params: {params_2_layers})\n"
-        else:
-            if target_p_avg < 50000: suggested_layers_str = str(max(1, int(target_p_avg / (input_dim + output_dim + 100)))) or "32"
-            elif target_p_avg < 250000: h = max(1,int(math.sqrt(target_p_avg/1.5))); suggested_layers_str=f"{h},{h//2}" if h>0 and h//2 >0 else "128,64"
-            else: h = max(1,int(math.sqrt(target_p_avg/2.0))); suggested_layers_str=f"{h},{h},{h//2}" if h>0 and h//2 >0 else "256,256,128"
-            logs += f"Fallback suggestion: {suggested_layers_str} (Verify params).\n"
-    if not suggested_layers_str: suggested_layers_str = "64"; logs += "Defaulting to '64'.\n"
-    return suggested_layers_str, logs
-
-def estimate_current_mlp_params(input_dim_str, hidden_layers_str, output_dim_str, task_type, current_logs=""):
-    logs = current_logs
-    try:
-        input_dim = int(input_dim_str); output_dim = int(output_dim_str)
-        if input_dim <= 0 or output_dim <= 0: return "Input/Output dims must be > 0", logs
-        
-        mlp_task_type = "classification" if task_type.endswith("Classification") else "regression"
-        temp_mlp = SimpleMLP(input_dim, hidden_layers_str, output_dim, task_type=mlp_task_type)
-        params = count_pytorch_parameters(temp_mlp); del temp_mlp
-        return f"{params:,}", logs
-    except Exception as e: logs += f"Error estimating MLP params: {e}\n"; return "Error", logs
-
-def estimate_cnn_params(img_h_str, img_w_str, num_classes_str, task_type, current_logs=""):
-    logs = current_logs
-    try:
-        img_h, img_w, num_classes_parsed = int(img_h_str), int(img_w_str), int(num_classes_str)
-        if not (img_h > 0 and img_w > 0 and num_classes_parsed > 0): return "Image dims/classes must be > 0", logs
-        
-        cnn_task_type = "classification" if task_type.endswith("Classification") else "regression"
-        temp_cnn = SimpleCNN(input_channels=1, img_size_wh=(img_h, img_w), num_classes=num_classes_parsed, task_type=cnn_task_type)
-        params = count_pytorch_parameters(temp_cnn); del temp_cnn
-        return f"{params:,}", logs
-    except Exception as e: logs += f"Error estimating CNN params: {traceback.format_exc()}\n"; return "Error", logs
-
-
-# --- Dataset and Preprocessing ---
-def generate_dataset_backend(task_type, n_samples_str, n_features_str,
-                             n_classes_or_informative_str, dataset_format,
-                             ai_suggest_ds_shape, target_param_range_str, model_type_selection,
-                             current_logs=""):
-    logs = current_logs + "\n--- Generating Dataset ---\n"
-    try:
-        n_samples = int(n_samples_str); n_features = int(n_features_str); n_classes_or_informative = int(n_classes_or_informative_str)
-    except ValueError: logs += "Invalid numbers for dataset generation.\n"; return None, "Error", logs, None
-
-    if ai_suggest_ds_shape:
-        n_samples_sugg, n_features_sugg, n_classes_or_informative_sugg = 5000, 10, 2
-        if task_type == "Tabular Regression": n_classes_or_informative_sugg = min(n_features_sugg // 2, 5) if n_features_sugg > 0 else 1
-        elif task_type == "Basic Image Classification": n_samples_sugg, n_features_sugg = 500, 0 
-        
-        is_nn = "Network" in model_type_selection
-        if is_nn and target_param_range_str in PARAM_RANGES:
-            min_p, max_p = PARAM_RANGES[target_param_range_str]; avg_p = (min_p + max_p) / 2
-            if avg_p > 200000: n_samples_sugg = min(MAX_GENERATED_ROWS, n_samples_sugg * 3); n_features_sugg = min(MAX_GENERATED_COLS, n_features_sugg * 2) if task_type.startswith("Tabular") else n_features_sugg
-            elif avg_p < 50000: n_samples_sugg = max(200, n_samples_sugg // 2); n_features_sugg = max(3, n_features_sugg // 2) if task_type.startswith("Tabular") else n_features_sugg
-        
-        n_samples, n_features, n_classes_or_informative = n_samples_sugg, n_features_sugg, n_classes_or_informative_sugg
-        logs += f"AI Suggested Dataset: Samples={n_samples}, Feats={n_features}, Classes/Informative={n_classes_or_informative}\n"
+        return self.network(x)
+
+# --- Dataset and Preprocessing Logic ---
+def generate_dataset_backend(task_type, n_samples, n_features, n_classes_or_informative, dataset_format):
+    """Generates synthetic data based on user specifications."""
+    logs = "\n--- Generating Dataset ---\n"
+    n_samples = max(10, min(int(n_samples), MAX_GENERATED_ROWS))
+    n_features = max(1, min(int(n_features), MAX_GENERATED_COLS))
+    n_classes_or_informative = int(n_classes_or_informative)
+    df = None
     
-    n_samples = max(10, min(n_samples, MAX_GENERATED_ROWS))
-    if task_type.startswith("Tabular"): n_features = max(1, min(n_features, MAX_GENERATED_COLS))
-    if n_samples > MAX_DATASET_ROWS_WARN: logs += f"Warning: Generating {n_samples} rows. May be slow.\n"
-
-    df = None; X_data=None; y_data=None
     try:
         if task_type == "Tabular Classification":
-            n_cls = max(2, n_classes_or_informative)
-            n_inf = max(1, min(n_features, n_classes_or_informative if n_classes_or_informative >= n_cls else n_features // 2))
-            if n_inf > n_features: n_inf = n_features
-            X_data, y_data = make_classification(n_samples=n_samples, n_features=n_features, n_informative=n_inf,
-                                       n_redundant=max(0,n_features - n_inf)//2, n_classes=n_cls, flip_y=0.05, random_state=42)
-            df = pd.DataFrame(X_data, columns=[f'feature_{i}' for i in range(n_features)]); df['target'] = y_data
+            X, y = make_classification(n_samples=n_samples, n_features=n_features, n_informative=max(1, n_features // 2),
+                                       n_redundant=0, n_classes=max(2, n_classes_or_informative), random_state=42)
+            df = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(n_features)])
+            df['target'] = y
         elif task_type == "Tabular Regression":
-            n_inf = max(1, min(n_features, n_classes_or_informative))
-            if n_inf > n_features: n_inf = n_features
-            X_data, y_data = make_regression(n_samples=n_samples, n_features=n_features, n_informative=n_inf, noise=10, random_state=42)
-            df = pd.DataFrame(X_data, columns=[f'feature_{i}' for i in range(n_features)]); df['target'] = y_data
-        elif task_type == "Basic Image Classification":
-            img_h, img_w = 28, 28 
-            num_pixels = img_h * img_w
-            X_data = np.random.randint(0, 256, size=(n_samples, num_pixels), dtype=np.uint8)
-            y_data = np.random.randint(0, max(2, n_classes_or_informative), n_samples)
-            df = pd.DataFrame(X_data, columns=[f'pixel_{i}' for i in range(num_pixels)]); df['target'] = y_data
-            logs += f"Generated {img_h}x{img_w} Image placeholder data.\n"
-        else: logs += f"Dataset generation for '{task_type}' not fully implemented.\n"; return None, "Task not implemented", logs, None
+            X, y = make_regression(n_samples=n_samples, n_features=n_features, 
+                                   n_informative=max(1, min(n_features, n_classes_or_informative)), noise=10, random_state=42)
+            df = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(n_features)])
+            df['target'] = y
         
-        logs += f"Generated data: {df.shape if df is not None else (X_data.shape, y_data.shape)}\n"
+        if df is None:
+            raise NotImplementedError(f"Dataset generation for '{task_type}' is not implemented.")
+            
+        logs += f"Generated data with shape: {df.shape}\n"
         file_path = get_temp_filepath("generated_dataset", dataset_format)
-        if df is not None: 
-            if dataset_format == ".csv": df.to_csv(file_path, index=False)
-            elif dataset_format == ".json": df.to_json(file_path, orient='records', lines=True)
-            elif dataset_format == ".parquet": df.to_parquet(file_path, index=False)
-            else: logs += f"Unsupported format {dataset_format}. Defaulting to CSV.\n"; file_path=get_temp_filepath("generated_dataset","csv"); df.to_csv(file_path, index=False)
-            logs += f"Dataset saved to {file_path}\n"
-            return df.head(), df, logs, file_path 
-        else:
-            logs += "Dataset generated as numpy arrays. Not saving to file from this function directly.\n"
-            return pd.DataFrame(X_data[:5] if X_data is not None else None), (X_data, y_data), logs, None
-
-    except Exception as e: error_msg=f"Error generating dataset: {traceback.format_exc()}"; logs+=error_msg+"\n"; return None, error_msg, logs, None
-
-def preprocess_tabular_data(df_or_X, y_if_X_is_numpy, target_column_name, task_type, current_logs=""):
-    logs = current_logs
-    if isinstance(df_or_X, pd.DataFrame):
-        df = df_or_X
-        if target_column_name not in df.columns: raise ValueError(f"Target column '{target_column_name}' not found.")
-        X_df = df.drop(target_column_name, axis=1)
-        y_series = df[target_column_name]
-    elif isinstance(df_or_X, np.ndarray) and y_if_X_is_numpy is not None:
-        X_df = pd.DataFrame(df_or_X, columns=[f'feature_{i}' for i in range(df_or_X.shape[1])])
-        y_series = pd.Series(y_if_X_is_numpy)
-    else: raise ValueError("Invalid input for preprocess_tabular_data.")
-
-    numerical_features = X_df.select_dtypes(include=np.number).columns.tolist()
-    categorical_features = X_df.select_dtypes(include='object').columns.tolist()
-    logs += f"Numerical: {numerical_features}, Categorical: {categorical_features}\n"
-
-    preprocessor = ColumnTransformer(transformers=[
-        ('num', Pipeline([('imputer', SimpleImputer(strategy='mean')), ('scaler', StandardScaler())]), numerical_features),
-        ('cat', Pipeline([('imputer', SimpleImputer(strategy='most_frequent')), ('onehot', OneHotEncoder(handle_unknown='ignore', sparse_output=False))]), categorical_features)
-    ], remainder='passthrough')
+        
+        if dataset_format == ".csv": df.to_csv(file_path, index=False)
+        elif dataset_format == ".json": df.to_json(file_path, orient='records', lines=True)
+        elif dataset_format == ".parquet": df.to_parquet(file_path, index=False)
+        
+        logs += f"Dataset saved to temporary file: {os.path.basename(file_path)}\n"
+        return df.head(), df, logs, file_path
 
-    X_processed_np = preprocessor.fit_transform(X_df)
-    processed_input_dim = X_processed_np.shape[1]
-    logs += f"Tabular data preprocessed. X shape: {X_processed_np.shape}, Processed input dim: {processed_input_dim}\n"
+    except Exception as e:
+        error_msg = f"Error generating dataset: {traceback.format_exc()}"
+        logs += error_msg + "\n"
+        return None, None, logs, None
 
-    if task_type.endswith("Classification"):
-        le = LabelEncoder()
-        y_processed_np = le.fit_transform(y_series)
-        num_classes = len(le.classes_)
-        logs += f"Target encoded. Classes: {num_classes} ({le.classes_})\n"
-        output_dim_nn = 1 if num_classes == 2 else num_classes
-    else: # Regression
-        y_processed_np = y_series.astype(float).values
-        num_classes = 1 
-        output_dim_nn = 1
+# --- Core Training Functions ---
+def train_model_sklearn(data_input, target_column, task_type, model_name, model_output_format, logs=""):
+    """Handles the entire Scikit-learn training and evaluation pipeline."""
+    logs += f"\n--- Training Scikit-learn Model: {model_name} ---\n"
     
-    return X_processed_np, y_processed_np, preprocessor, logs, processed_input_dim, output_dim_nn
-
-
-# --- Training Functions ---
-def train_model_sklearn(data_input_obj, target_column, task_type, model_name, model_output_format, current_logs=""):
-    logs = current_logs + f"\n--- Training Scikit-learn Model: {model_name} ---\n"
-    model_path_out, metrics_out, model_params_out = None, "Training failed.", "N/A"
-
-    df = None
-    if isinstance(data_input_obj, str): 
-        try:
-            if data_input_obj.endswith('.csv'): df = pd.read_csv(data_input_obj)
-            elif data_input_obj.endswith('.json'): df = pd.read_json(data_input_obj, lines=True)
-            elif data_input_obj.endswith('.parquet'): df = pd.read_parquet(data_input_obj)
-            else: logs += f"Unsupported file: {data_input_obj}\n"; return logs, "Error: Unsupported file.", None, "N/A"
-        except Exception as e: logs += f"Error reading {data_input_obj}: {e}\n"; return logs, f"Error reading: {e}", None, "N/A"
-    elif isinstance(data_input_obj, pd.DataFrame): df = data_input_obj
-    else: logs += "Invalid data for training.\n"; return logs, "Error: Invalid data.", None, "N/A"
-
-    if not target_column or target_column not in df.columns:
-        logs += f"Target column '{target_column}' not provided or not found.\n"; return logs, f"Error: Target '{target_column}' not found/provided.", None, "N/A"
-
     try:
-        X_processed_np, y_processed_np, preprocessor, logs, _, _ = preprocess_tabular_data(df, None, target_column, task_type, logs)
-    except ValueError as e: logs += f"Preprocessing error: {e}\n"; return logs, f"Error: {e}", None, "N/A"
-
-    X_train, X_test, y_train, y_test = train_test_split(X_processed_np, y_processed_np, test_size=0.2, random_state=42)
-    logs += f"Train/Test split. Train: {X_train.shape}, Test: {X_test.shape}\n"
-
-    model = None
-    if task_type == "Tabular Classification":
-        if model_name == "Logistic Regression": model = LogisticRegression(max_iter=1000, random_state=42)
-        elif model_name == "Random Forest Classifier": model = RandomForestClassifier(random_state=42)
-        elif model_name == "Support Vector Machine (SVM) Classifier": model = SVC(random_state=42, probability=True) 
-    elif task_type == "Tabular Regression":
-        if model_name == "Linear Regression": model = LinearRegression()
-        elif model_name == "Random Forest Regressor": model = RandomForestRegressor(random_state=42)
-        elif model_name == "Support Vector Machine (SVR) Regressor": model = SVR()
-    if model is None: logs += f"Model {model_name} or task {task_type} not supported.\n"; return logs, "Model/Task Error", None, "N/A"
-
-    try:
-        logs += f"Starting training for {model_name}...\n"; start_time = time.time()
-        model.fit(X_train, y_train)
+        if isinstance(data_input, str): # Is a filepath
+            if data_input.endswith('.csv'): df = pd.read_csv(data_input)
+            else: raise ValueError("Unsupported file type for upload.")
+        else: # Is a DataFrame from generation
+            df = data_input
+
+        if target_column not in df.columns:
+            raise ValueError(f"Target column '{target_column}' not found.")
+
+        # Preprocessing
+        X = df.drop(columns=[target_column])
+        y = df[target_column]
+        numeric_features = X.select_dtypes(include=np.number).columns
+        categorical_features = X.select_dtypes(include='object').columns
+        
+        preprocessor = ColumnTransformer(transformers=[
+            ('num', Pipeline([('imputer', SimpleImputer(strategy='mean')), ('scaler', StandardScaler())]), numeric_features),
+            ('cat', Pipeline([('imputer', SimpleImputer(strategy='most_frequent')), ('onehot', OneHotEncoder(handle_unknown='ignore'))]), categorical_features)
+        ])
+        
+        # Model Selection
+        if task_type == "Tabular Classification":
+            y = LabelEncoder().fit_transform(y)
+            models = {
+                "Logistic Regression": LogisticRegression(max_iter=1000, random_state=42),
+                "Random Forest Classifier": RandomForestClassifier(random_state=42),
+                "Support Vector Machine (SVM) Classifier": SVC(random_state=42, probability=True)
+            }
+        else: # Regression
+            models = {
+                "Linear Regression": LinearRegression(),
+                "Random Forest Regressor": RandomForestRegressor(random_state=42),
+                "Support Vector Machine (SVR) Regressor": SVR()
+            }
+        model = models[model_name]
+        
+        # Create full pipeline
+        pipeline = Pipeline([('preprocessor', preprocessor), ('model', model)])
+        
+        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+        logs += f"Data split into training ({X_train.shape}) and testing ({X_test.shape}) sets.\n"
+        
+        # Training
+        start_time = time.time()
+        pipeline.fit(X_train, y_train)
         logs += f"Training completed in {time.time() - start_time:.2f}s.\n"
-        model_params_out = str(count_sklearn_parameters(model))
-        logs += f"Est. Model Params: {model_params_out}\n"
-        y_pred = model.predict(X_test)
-
+        
+        # Evaluation
+        y_pred = pipeline.predict(X_test)
         if task_type == "Tabular Classification":
-            acc = accuracy_score(y_test, y_pred); report = classification_report(y_test, y_pred, zero_division=0)
-            metrics_out = f"Accuracy: {acc:.4f}\n\nClassification Report:\n{report}"
-        elif task_type == "Tabular Regression":
-            mse = mean_squared_error(y_test, y_pred); r2 = r2_score(y_test, y_pred)
-            metrics_out = f"Mean Squared Error: {mse:.4f}\nR2 Score: {r2:.4f}"
-        logs += "\n--- Evaluation Metrics ---\n" + metrics_out + "\n"
-        full_pipeline_for_saving = Pipeline([('preprocessor', preprocessor), ('model', model)])
+            acc = accuracy_score(y_test, y_pred)
+            report = classification_report(y_test, y_pred, zero_division=0)
+            metrics = f"Accuracy: {acc:.4f}\n\nClassification Report:\n{report}"
+        else:
+            mse = mean_squared_error(y_test, y_pred)
+            r2 = r2_score(y_test, y_pred)
+            metrics = f"Mean Squared Error: {mse:.4f}\nR² Score: {r2:.4f}"
+        logs += "\n--- Evaluation Metrics ---\n" + metrics + "\n"
+        
+        # Model Saving
         model_filename_base = f"sklearn_{model_name.replace(' ', '_').lower()}"
-
         if model_output_format == ".pkl (Scikit-learn)":
-            model_path_out = get_temp_filepath(model_filename_base, "pkl")
-            joblib.dump(full_pipeline_for_saving, model_path_out)
-            logs += f"Model (with preprocessor) saved to {model_path_out} as PKL.\n"
+            model_path = get_temp_filepath(model_filename_base, "pkl")
+            joblib.dump(pipeline, model_path)
+            logs += f"Model pipeline saved to {os.path.basename(model_path)} as PKL.\n"
         elif model_output_format == ".onnx (ONNX)":
-            model_path_out = get_temp_filepath(model_filename_base, "onnx")
-            raw_X_for_types_df = df.drop(target_column, axis=1).infer_objects()
-            onnx_initial_types = []
-            for col_name in raw_X_for_types_df.columns:
-                col_dtype = raw_X_for_types_df[col_name].dtype
-                if pd.api.types.is_numeric_dtype(col_dtype):
-                    onnx_initial_types.append((col_name, FloatTensorType([None, 1])))
-                elif pd.api.types.is_string_dtype(col_dtype) or col_dtype == 'object':
-                     onnx_initial_types.append((col_name, StringTensorType([None, 1])))
+            model_path = get_temp_filepath(model_filename_base, "onnx")
+            initial_types = []
+            for col_name in X.columns:
+                if pd.api.types.is_numeric_dtype(X[col_name].dtype):
+                    initial_types.append((col_name, FloatTensorType([None, 1])))
                 else:
-                    logs += f"Warning: Unsupported dtype {col_dtype} for {col_name} in ONNX. Defaulting to Float.\n"
-                    onnx_initial_types.append((col_name, FloatTensorType([None, 1])))
+                    initial_types.append((col_name, StringTensorType([None, 1])))
             
-            if not onnx_initial_types: raise ValueError("ONNX initial types failed: No valid columns found.")
-            try:
-                options = {id(full_pipeline_for_saving): {'zipmap': False}} if task_type.endswith("Classification") else {}
-                onnx_model = convert_sklearn(full_pipeline_for_saving, initial_types=onnx_initial_types, target_opset=12, options=options)
-                with open(model_path_out, "wb") as f: f.write(onnx_model.SerializeToString())
-                logs += f"Model saved to {model_path_out} as ONNX.\n"
-                sess = rt.InferenceSession(model_path_out, providers=rt.get_available_providers())
-                logs += f"ONNX model loaded successfully with ONNX Runtime.\n"
-            except Exception as onnx_e: logs += f"ONNX Error: {traceback.format_exc()}\n"; model_path_out=None; metrics_out+="\nONNX EXPORT FAILED."
-        else:
-            logs += f"Unsupported format '{model_output_format}'. Saving as .pkl\n"
-            model_path_out = get_temp_filepath(model_filename_base, "pkl")
-            joblib.dump(full_pipeline_for_saving, model_path_out)
-    except Exception as e: error_msg=f"Sklearn training error: {traceback.format_exc()}"; logs+=error_msg+"\n"; metrics_out=error_msg
-    return logs, metrics_out, model_path_out, model_params_out
-
-def train_model_pytorch(data_input_obj, target_column, task_type, model_type_pt,
-                        mlp_hidden_layers_str, mlp_activation,
-                        epochs_str, batch_size_str, lr_str,
-                        model_output_format, current_logs=""):
-    logs = current_logs + f"\n--- Training PyTorch Model: {model_type_pt} ---\n"
-    model_path_out, metrics_out, model_params_out, plot_out = None, "Training failed.", "N/A", None
-
-    df_for_pytorch, X_numpy_for_pytorch, y_numpy_for_pytorch = None, None, None
-    if isinstance(data_input_obj, str): 
-        try:
-            if data_input_obj.endswith('.csv'): df_for_pytorch = pd.read_csv(data_input_obj)
-            elif data_input_obj.endswith('.json'): df_for_pytorch = pd.read_json(data_input_obj, lines=True)
-            elif data_input_obj.endswith('.parquet'): df_for_pytorch = pd.read_parquet(data_input_obj)
-            else: logs += f"Unsupported file: {data_input_obj}\n"; return logs, "Error", None, "N/A", None
-        except Exception as e: logs += f"Error reading {data_input_obj}: {e}\n"; return logs, f"Error: {e}", None, "N/A", None
-    elif isinstance(data_input_obj, pd.DataFrame): df_for_pytorch = data_input_obj
-    elif isinstance(data_input_obj, tuple): X_numpy_for_pytorch, y_numpy_for_pytorch = data_input_obj
-    else: logs += "Invalid data for PyTorch training.\n"; return logs, "Error", None, "N/A", None
-
-    try:
-        epochs = int(epochs_str); batch_size = int(batch_size_str); lr = float(lr_str)
-        if not (epochs > 0 and batch_size > 0 and lr > 0): raise ValueError("Params must be >0.")
-    except ValueError as e: logs += f"Invalid training params: {e}\n"; return logs, f"Error: {e}", None, "N/A", None
-
-    processed_input_dim_actual, nn_output_dim_actual, preprocessor_pipeline = -1, -1, None
-    X_processed_np, y_processed_np = None, None
-
-    if model_type_pt == "Simple Neural Network (MLP)":
-        if not task_type.startswith("Tabular"): logs += "MLP requires Tabular task.\n"; return logs, "MLP Task Error", None, "N/A", None
-        if not target_column and df_for_pytorch is not None: logs += "Target column needed for MLP with DataFrame.\n"; return logs, "MLP Target Error", None, "N/A", None
-        try:
-            data_arg1 = df_for_pytorch if df_for_pytorch is not None else X_numpy_for_pytorch
-            data_arg2 = y_numpy_for_pytorch if df_for_pytorch is None else None
-            current_target_col = target_column if df_for_pytorch is not None else "target"
-            X_processed_np, y_processed_np, preprocessor_pipeline, logs, processed_input_dim_actual, nn_output_dim_actual = \
-                preprocess_tabular_data(data_arg1, data_arg2, current_target_col, task_type, logs)
-        except ValueError as e: logs+=f"MLP Preprocessing error: {e}\n"; return logs,f"Error: {e}",None,"N/A",None
-    elif model_type_pt == "Simple Convolutional Network (CNN)":
-        X_raw, y_raw = (df_for_pytorch.drop(target_column, axis=1).values, df_for_pytorch[target_column].values) if df_for_pytorch is not None else (X_numpy_for_pytorch, y_numpy_for_pytorch)
-        if X_raw is None: logs += "No valid data for CNN.\n"; return logs, "CNN Data Error", None, "N/A", None
-        le = LabelEncoder(); y_processed_np = le.fit_transform(y_raw)
-        num_classes = len(le.classes_)
-        nn_output_dim_actual = 1 if num_classes == 2 else num_classes
-
-        pixels_per_sample = X_raw.shape[1]; img_h, img_w, input_channels = 28,28,1 
-        img_dim_approx = int(math.sqrt(pixels_per_sample))
-        if img_dim_approx * img_dim_approx == pixels_per_sample: img_h, img_w = img_dim_approx, img_dim_approx
-        else: logs += f"Warning: Cannot infer square image from {pixels_per_sample} pixels. Defaulting to 28x28.\n"
-        
-        X_processed_np = X_raw.reshape(-1, input_channels, img_h, img_w).astype(np.float32) / 255.0
-        processed_input_dim_actual = (input_channels, img_h, img_w) 
-        logs += f"CNN Data: X reshaped to {X_processed_np.shape}, y: {y_processed_np.shape}\n"
-    else: logs += f"Unknown PyTorch model: {model_type_pt}\n"; return logs, "Unknown PyTorch model", None, "N/A", None
-
-    X_train, X_test, y_train, y_test = train_test_split(X_processed_np, y_processed_np, test_size=0.2, random_state=42)
-    logs += f"PyTorch Train/Test split. Train: {X_train.shape}, Test: {X_test.shape}\n"
-
-    y_train_dtype = torch.float32 if (nn_output_dim_actual == 1 and not task_type.endswith("Classification")) else (torch.float32 if nn_output_dim_actual == 1 else torch.long)
-    X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
-    y_train_tensor = torch.tensor(y_train, dtype=y_train_dtype)
-    X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
-    y_test_tensor = torch.tensor(y_test, dtype=y_train_dtype)
-
-    if nn_output_dim_actual == 1: # For BCELoss and MSELoss, target needs to be [N, 1]
-        y_train_tensor = y_train_tensor.unsqueeze(1)
-        y_test_tensor = y_test_tensor.unsqueeze(1)
-
-    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
-    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
-
-    pytorch_model = None
-    try:
-        is_classification = task_type.endswith("Classification")
-        if model_type_pt == "Simple Neural Network (MLP)":
-            pytorch_model = SimpleMLP(input_dim=processed_input_dim_actual, hidden_layers_str=mlp_hidden_layers_str,
-                                    output_dim=nn_output_dim_actual, activation_fn_str=mlp_activation,
-                                    task_type="classification" if is_classification else "regression")
-        elif model_type_pt == "Simple Convolutional Network (CNN)":
-            channels, h, w = processed_input_dim_actual
-            pytorch_model = SimpleCNN(input_channels=channels, img_size_wh=(h,w), num_classes=nn_output_dim_actual,
-                                      task_type="classification" if is_classification else "regression")
-    except Exception as model_e: logs += f"Error creating PyTorch model: {traceback.format_exc()}\n"; return logs, f"Model Creation Error: {model_e}", None, "N/A", None
-
-    if pytorch_model is None: logs += "Failed to instantiate PyTorch model.\n"; return logs, "Model instantiate fail", None, "N/A", None
-    model_params_val = count_pytorch_parameters(pytorch_model); model_params_out = f"{model_params_val:,}"
-    logs += f"PyTorch Model: {model_params_out} params.\n"
-
-    criterion = nn.BCELoss() if (is_classification and nn_output_dim_actual == 1) else (nn.CrossEntropyLoss() if is_classification else nn.MSELoss())
-    optimizer = optim.Adam(pytorch_model.parameters(), lr=lr)
-
-    logs += f"Starting PyTorch training for {epochs} epochs...\n"; start_time = time.time()
-    epoch_losses = []
-    pytorch_model.train()
-    for epoch in range(epochs):
-        epoch_loss_sum = 0.0
-        for batch_X, batch_y in train_dataloader:
-            optimizer.zero_grad(); outputs = pytorch_model(batch_X)
-            loss = criterion(outputs, batch_y); loss.backward(); optimizer.step()
-            epoch_loss_sum += loss.item()
-        avg_epoch_loss = epoch_loss_sum / len(train_dataloader) if len(train_dataloader) > 0 else 0
-        epoch_losses.append(avg_epoch_loss)
-        if (epoch + 1) % max(1, epochs // 10) == 0 or epoch == epochs - 1: # Log ~10 times
-            logs += f"Epoch {epoch+1}/{epochs}, Avg Training Loss: {avg_epoch_loss:.4f}\n"
-    
-    logs += f"PyTorch training completed in {time.time() - start_time:.2f}s.\n"
-    pytorch_model.eval()
-    with torch.no_grad():
-        test_outputs = pytorch_model(X_test_tensor)
-        if is_classification:
-            predicted = (test_outputs > 0.5).float() if nn_output_dim_actual == 1 else torch.max(test_outputs.data, 1)[1]
-            acc = accuracy_score(y_test_tensor.cpu().numpy(), predicted.cpu().numpy())
-            report = classification_report(y_test_tensor.cpu().numpy(), predicted.cpu().numpy(), zero_division=0)
-            metrics_out = f"Final Avg Training Loss: {avg_epoch_loss:.4f}\n\n--- Test Set Evaluation ---\nAccuracy: {acc:.4f}\n\nClassification Report:\n{report}"
-        else: # Regression
-            mse = mean_squared_error(y_test_tensor.cpu().numpy(), test_outputs.cpu().numpy())
-            r2 = r2_score(y_test_tensor.cpu().numpy(), test_outputs.cpu().numpy())
-            metrics_out = f"Final Avg Training Loss: {avg_epoch_loss:.4f}\n\n--- Test Set Evaluation ---\nMean Squared Error: {mse:.4f}\nR2 Score: {r2:.4f}"
-    logs += "\n--- PyTorch Metrics ---\n" + metrics_out + "\n"
-
-    if epoch_losses:
-        fig, ax = plt.subplots(); ax.plot(range(1, epochs + 1), epoch_losses, marker='o')
-        ax.set_xlabel("Epoch"); ax.set_ylabel("Average Loss"); ax.set_title("Training Loss Curve")
-        plot_out = fig; logs += "Loss curve generated.\n"
-
-    model_filename_base = f"pytorch_{model_type_pt.replace(' ', '_').lower()}"
-    if model_output_format == ".pt (PyTorch)":
-        model_path_out = get_temp_filepath(model_filename_base, "pt")
-        save_obj = {'model_state_dict': pytorch_model.state_dict(), 'output_dim_nn': nn_output_dim_actual, 'task_type': task_type}
-        if model_type_pt == "Simple Neural Network (MLP)" and preprocessor_pipeline:
-            save_obj.update({'preprocessor': preprocessor_pipeline, 'input_dim_processed': processed_input_dim_actual, 'hidden_layers_str': mlp_hidden_layers_str, 'activation_fn': mlp_activation})
-        elif model_type_pt == "Simple Convolutional Network (CNN)":
-            c,h,w = processed_input_dim_actual; save_obj.update({'input_channels':c, 'img_h':h, 'img_w':w})
-        torch.save(save_obj, model_path_out)
-        logs += f"PyTorch model saved to {model_path_out}\n"
-    else: # Fallback
-        logs += f"Unsupported format '{model_output_format}'.\n"
-    return logs, metrics_out, model_path_out, model_params_out, plot_out
-
-
-# --- Main Training Wrapper Function ---
-def train_model_wrapper(data_input_obj, target_column, task_type, model_family, 
-                        model_specific_choice, 
-                        mlp_hidden_layers, mlp_activation,
-                        epochs, batch_size, learning_rate,
-                        model_output_format, current_logs):
-    
-    logs = current_logs + "\n--- Kicking off Training ---\n"
-    if data_input_obj is None:
-        logs += "ERROR: No dataset has been generated or uploaded. Please go to Tab 2.\n"
-        return logs, "Error: No dataset available.", None, "N/A", None
+            onnx_model = convert_sklearn(pipeline, initial_types=initial_types, target_opset=12)
+            with open(model_path, "wb") as f: f.write(onnx_model.SerializeToString())
+            logs += f"Model pipeline saved to {os.path.basename(model_path)} as ONNX.\n"
+            
+            if ONNX_RUNTIME_AVAILABLE:
+                sess = rt.InferenceSession(model_path)
+                logs += "ONNX model successfully loaded and validated with onnxruntime.\n"
+            else:
+                logs += "ONNX model validation skipped because onnxruntime is not available in this environment.\n"
+                
+        return logs, metrics, model_path
 
-    try:
-        if model_family == "Scikit-learn (Classical ML)":
-            logs, metrics, model_path, param_count = train_model_sklearn(
-                data_input_obj=data_input_obj,
-                target_column=target_column,
-                task_type=task_type,
-                model_name=model_specific_choice,
-                model_output_format=model_output_format,
-                current_logs=logs
-            )
-            return logs, metrics, model_path, param_count, None
-        elif model_family == "PyTorch (Neural Networks)":
-            logs, metrics, model_path, param_count, loss_plot = train_model_pytorch(
-                data_input_obj=data_input_obj,
-                target_column=target_column,
-                task_type=task_type,
-                model_type_pt=model_specific_choice,
-                mlp_hidden_layers_str=mlp_hidden_layers,
-                mlp_activation=mlp_activation,
-                epochs_str=str(int(epochs)),
-                batch_size_str=str(int(batch_size)),
-                lr_str=str(learning_rate),
-                model_output_format=model_output_format,
-                current_logs=logs
-            )
-            return logs, metrics, model_path, param_count, loss_plot
-        else:
-            logs += f"Unknown model family: {model_family}\n"
-            return logs, "Error: Unknown model family.", None, "N/A", None
     except Exception as e:
-        error_msg = f"An unexpected error occurred in the training wrapper: {traceback.format_exc()}"
+        error_msg = f"Scikit-learn training failed: {traceback.format_exc()}"
         logs += error_msg + "\n"
-        return logs, error_msg, None, "N/A", None
-
+        return logs, error_msg, None
 
-# --- Gradio UI Definition ---
-TASK_CHOICES = ["Tabular Classification", "Tabular Regression", "Basic Image Classification"]
-MODEL_FAMILIES = ["Scikit-learn (Classical ML)", "PyTorch (Neural Networks)"]
-SKLEARN_MODELS_CLASSIFICATION = ["Logistic Regression", "Random Forest Classifier", "Support Vector Machine (SVM) Classifier"]
-SKLEARN_MODELS_REGRESSION = ["Linear Regression", "Random Forest Regressor", "Support Vector Machine (SVR) Regressor"]
-PYTORCH_MODELS = ["Simple Neural Network (MLP)", "Simple Convolutional Network (CNN)"]
-DATASET_FORMATS = [".csv", ".json", ".parquet"]
-MODEL_OUTPUT_FORMATS_SKLEARN = [".pkl (Scikit-learn)", ".onnx (ONNX)"]
-MODEL_OUTPUT_FORMATS_PYTORCH = [".pt (PyTorch)"]
-MLP_ACTIVATIONS = ["relu", "tanh", "sigmoid"]
-CLONE_GUIDE_TEXT = """
-## How to Clone & Upgrade This Space for More Power:
-1.  **Clone this Space:** Click the '...' menu at the top-right and choose 'Duplicate this Space'.
-2.  **Choose Hardware:** On the duplication screen, select a more powerful hardware option, like a "CPU upgrade" or a "T4 Small" GPU.
-3.  **Enjoy Faster Training:** Your private, upgraded version of TrainAI will now train models significantly faster!
-"""
+# --- Main Training Dispatcher ---
+def train_model_wrapper(data_input, target_column, task_type, model_family, model_specific, 
+                        model_output_format, logs):
+    """A wrapper to call the correct training function based on user choices."""
+    if data_input is None:
+        logs += "ERROR: No dataset has been generated or uploaded. Please go to Tab 2.\n"
+        return logs, "Error: No dataset available.", None, None
 
-# Determine initial choices for model_specific_dd based on default task_type and model_family
-_initial_task_default = TASK_CHOICES[0]
-_initial_family_default = MODEL_FAMILIES[0]
-if _initial_family_default == "Scikit-learn (Classical ML)":
-    initial_model_choices_for_specific_dd = SKLEARN_MODELS_CLASSIFICATION
-elif _initial_family_default == "PyTorch (Neural Networks)":
-    initial_model_choices_for_specific_dd = [PYTORCH_MODELS[0]]
-initial_model_value_for_specific_dd = initial_model_choices_for_specific_dd[0]
+    if model_family == "Scikit-learn (Classical ML)":
+        logs, metrics, model_path = train_model_sklearn(data_input, target_column, task_type, model_specific, model_output_format, logs)
+        return logs, metrics, model_path, None # No plot for sklearn
+    
+    # Placeholder for PyTorch integration if added back
+    elif model_family == "PyTorch (Neural Networks)":
+        logs += "PyTorch training is not fully integrated in this version yet.\n"
+        return logs, "PyTorch not available.", None, None
+        
+    else:
+        logs += f"Unknown model family: {model_family}\n"
+        return logs, "Error: Unknown model family.", None, None
 
+# --- Gradio UI Definition ---
 def update_model_options(task_choice, model_family_choice):
-    choices, value = [], None
+    """Dynamically updates the available models based on task and family."""
+    choices = []
     if model_family_choice == "Scikit-learn (Classical ML)":
-        if task_choice == "Tabular Classification": choices = SKLEARN_MODELS_CLASSIFICATION
-        elif task_choice == "Tabular Regression": choices = SKLEARN_MODELS_REGRESSION
-    elif model_family_choice == "PyTorch (Neural Networks)":
-        if task_choice.startswith("Tabular"): choices = [PYTORCH_MODELS[0]] # MLP
-        elif task_choice == "Basic Image Classification": choices = [PYTORCH_MODELS[1]] # CNN
+        if task_choice == "Tabular Classification":
+            choices = ["Logistic Regression", "Random Forest Classifier", "Support Vector Machine (SVM) Classifier"]
+        elif task_choice == "Tabular Regression":
+            choices = ["Linear Regression", "Random Forest Regressor", "Support Vector Machine (SVR) Regressor"]
+    # Add PyTorch options here if needed
+    
     value = choices[0] if choices else None
     return gr.update(choices=choices, value=value, visible=bool(choices))
 
-def update_pytorch_specific_options_visibility(model_family_choice, specific_pytorch_model):
-    is_pytorch = model_family_choice == "PyTorch (Neural Networks)"
-    is_mlp = is_pytorch and (specific_pytorch_model == "Simple Neural Network (MLP)")
-    is_cnn = is_pytorch and (specific_pytorch_model == "Simple Convolutional Network (CNN)")
-    return gr.update(visible=is_pytorch), gr.update(visible=is_mlp), gr.update(visible=is_cnn)
-
 def update_model_output_formats(model_family_choice):
-    if model_family_choice == "Scikit-learn (Classical ML)": return gr.update(choices=MODEL_OUTPUT_FORMATS_SKLEARN, value=MODEL_OUTPUT_FORMATS_SKLEARN[0])
-    if model_family_choice == "PyTorch (Neural Networks)": return gr.update(choices=MODEL_OUTPUT_FORMATS_PYTORCH, value=MODEL_OUTPUT_FORMATS_PYTORCH[0])
-    return gr.update(choices=[], value=None)
-
-css = """.gradio-container { font-family: 'IBM Plex Sans', sans-serif; } footer {display:none !important}"""
+    """Updates the output format options based on the model family."""
+    formats = []
+    if model_family_choice == "Scikit-learn (Classical ML)":
+        formats = [".pkl (Scikit-learn)", ".onnx (ONNX)"]
+    # Add PyTorch formats here
+    
+    value = formats[0] if formats else None
+    return gr.update(choices=formats, value=value)
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"), css=css) as demo:
+# The Gradio App Layout
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange")) as demo:
     gr.Markdown("# 🧠 TrainAI ⚙️")
-    gr.Markdown("Create, train, and download AI models. Optimized for CPU - expect longer training for complex models.")
-    generated_data_state = gr.State(None) 
-    current_logs_state = gr.State("")
-
+    gr.Markdown("A simple interface to create, train, and download machine learning models.")
+    
+    # State variables to hold data between interactions
+    generated_data_state = gr.State(None)
+    
     with gr.Tabs():
         with gr.TabItem("1. Define Task & Model"):
             with gr.Row():
-                task_type_dd = gr.Dropdown(TASK_CHOICES, label="Select Task Type", value=_initial_task_default)
-                model_family_dd = gr.Dropdown(MODEL_FAMILIES, label="Select Model Family", value=_initial_family_default)
-            model_specific_dd = gr.Dropdown(label="Select Specific Model", choices=initial_model_choices_for_specific_dd, value=initial_model_value_for_specific_dd, interactive=True)
+                task_type_dd = gr.Dropdown(["Tabular Classification", "Tabular Regression"], label="Select Task Type", value="Tabular Classification")
+                model_family_dd = gr.Dropdown(["Scikit-learn (Classical ML)"], label="Select Model Family", value="Scikit-learn (Classical ML)")
             
-            with gr.Group(visible=(_initial_family_default == "PyTorch (Neural Networks)")) as pt_options_group:
-                pytorch_param_range_dd = gr.Dropdown(list(PARAM_RANGES.keys()), label="Target Parameter Range (for NNs)", 
-                                                    info="Guides NN architecture suggestions. Training >250k params on CPU is slow.",
-                                                    value=list(PARAM_RANGES.keys())[1])
-                with gr.Group(visible=(initial_model_value_for_specific_dd == PYTORCH_MODELS[0])) as pt_mlp_specific_group:
-                    gr.Markdown("#### MLP Configuration")
-                    pt_mlp_hidden_layers_txt = gr.Textbox(label="Hidden Layer Sizes (comma-separated, e.g., 128,64)", value="64,32")
-                    pt_mlp_activation_dd = gr.Dropdown(MLP_ACTIVATIONS, label="Activation Function", value="relu")
-                    with gr.Row():
-                        pt_mlp_suggest_btn = gr.Button("Suggest MLP Layers")
-                        pt_mlp_estimate_params_btn = gr.Button("Estimate Current MLP Params")
-                    pt_mlp_param_count_txt = gr.Textbox(label="Estimated MLP Parameters", interactive=False)
-                with gr.Group(visible=(initial_model_value_for_specific_dd == PYTORCH_MODELS[1])) as pt_cnn_specific_group:
-                    gr.Markdown("#### CNN Configuration (Simplified)")
-                    gr.Markdown("SimpleCNN uses a fixed structure. Params depend on image size/classes from data.")
-                    pt_cnn_estimate_params_btn = gr.Button("Estimate CNN Params (needs Data Info)")
-                    pt_cnn_param_count_txt = gr.Textbox(label="Estimated CNN Parameters", interactive=False)
+            model_specific_dd = gr.Dropdown(label="Select Specific Model", choices=["Logistic Regression", "Random Forest Classifier", "Support Vector Machine (SVM) Classifier"], value="Logistic Regression", interactive=True)
 
         with gr.TabItem("2. Configure Dataset"):
-            dataset_source_rb = gr.Radio(["Generate new dataset", "Upload my own dataset (CSV, JSON, Parquet)"], 
-                                         label="Dataset Source", value="Generate new dataset")
-            with gr.Group(visible=True) as generate_dataset_group:
-                gr.Markdown("#### Generate Synthetic Dataset")
-                with gr.Row():
-                    ds_gen_samples_num = gr.Number(label="# Samples", value=1000, minimum=10, step=100)
-                    ds_gen_features_num = gr.Number(label="# Features (Tabular)", value=10, minimum=1, step=1)
-                ds_gen_classes_informative_num = gr.Number(label="Classes (Classif) / Informative Feats (Regr)", value=2, minimum=1, step=1)
-                ds_gen_ai_suggest_cb = gr.Checkbox(label="Let AI suggest dataset shape?", value=False)
-                ds_gen_format_dd = gr.Dropdown(DATASET_FORMATS, label="Generated Dataset Format", value=".csv")
-                generate_dataset_btn = gr.Button("Generate & Preview Dataset", variant="secondary")
-            with gr.Group(visible=False) as upload_dataset_group:
-                gr.Markdown("#### Upload Dataset")
-                ds_upload_file = gr.File(label="Upload dataset", file_types=[".csv", ".json", ".parquet"])
-            target_column_name_txt = gr.Textbox(label="Target Column Name (Case-Sensitive!)", placeholder="e.g., 'target' or 'label'", value="target")
+            with gr.Row():
+                ds_gen_samples_num = gr.Number(label="# Samples", value=1000, minimum=10, step=100)
+                ds_gen_features_num = gr.Number(label="# Features", value=10, minimum=1, step=1)
+                ds_gen_classes_num = gr.Number(label="Classes (Classif) / Informative (Regr)", value=2, minimum=1, step=1)
+            ds_gen_format_dd = gr.Dropdown([".csv", ".json", ".parquet"], label="Generated Dataset Format", value=".csv")
+            generate_dataset_btn = gr.Button("Generate & Preview Dataset", variant="secondary")
+            
+            target_column_name_txt = gr.Textbox(label="Target Column Name", value="target", interactive=True)
             dataset_preview_df = gr.DataFrame(label="Dataset Preview (First 5 Rows)", interactive=False, height=200)
             generated_dataset_download_file = gr.File(label="Download Generated Dataset", interactive=False)
 
         with gr.TabItem("3. Train Model & Get Results"):
-            gr.Markdown("Ensure Model and Dataset are configured. For NNs, set Epochs/Batch/LR.")
-            with gr.Row():
-                train_epochs_num = gr.Number(label="Epochs (NNs)", value=10, minimum=1, step=1)
-                train_batch_size_num = gr.Number(label="Batch Size (NNs)", value=32, minimum=1, step=1)
-                train_learning_rate_num = gr.Number(label="Learning Rate (NNs)", value=0.001, minimum=1e-6, format="%.6f")
-            model_output_format_dd = gr.Dropdown(label="Select Model Output Format", choices=MODEL_OUTPUT_FORMATS_SKLEARN, value=MODEL_OUTPUT_FORMATS_SKLEARN[0])
+            model_output_format_dd = gr.Dropdown(label="Select Model Output Format", choices=[".pkl (Scikit-learn)", ".onnx (ONNX)"], value=".pkl (Scikit-learn)")
             train_model_btn = gr.Button("🚀 Train Model", variant="primary")
             gr.Markdown("---")
             gr.Markdown("### Training Progress & Results")
             training_log_txt = gr.Textbox(label="Training Log & Status", lines=15, interactive=False, max_lines=50)
-            model_param_count_output_txt = gr.Textbox(label="Actual Trained Model Parameters", interactive=False)
             evaluation_metrics_txt = gr.Textbox(label="Evaluation Metrics", lines=7, interactive=False)
-            loss_plot_img = gr.Plot(label="Training Loss Curve (PyTorch NNs)")
             download_trained_model_file = gr.File(label="Download Trained Model", interactive=False)
+            loss_plot_img = gr.Plot(label="Training Loss Curve (PyTorch only)", visible=False) # Hide for now
 
-        with gr.TabItem("ℹ️ Guide & Info"):
-            gr.Markdown(CLONE_GUIDE_TEXT)
-            
     # --- Event Handlers ---
+    
+    # Update model choices when task or family changes
     task_type_dd.change(fn=update_model_options, inputs=[task_type_dd, model_family_dd], outputs=model_specific_dd)
     model_family_dd.change(fn=update_model_options, inputs=[task_type_dd, model_family_dd], outputs=model_specific_dd)
     
-    model_family_dd.change(fn=update_pytorch_specific_options_visibility, inputs=[model_family_dd, model_specific_dd], outputs=[pt_options_group, pt_mlp_specific_group, pt_cnn_specific_group])
-    model_specific_dd.change(fn=update_pytorch_specific_options_visibility, inputs=[model_family_dd, model_specific_dd], outputs=[pt_options_group, pt_mlp_specific_group, pt_cnn_specific_group])
-
-    def get_data_dims_for_nn_suggestion(preview_df, target_col, task, logs_in):
-        logs = logs_in
-        input_dim_est, output_dim_est = 10, (2 if task.endswith("Classification") else 1) # Defaults
-        img_h_est, img_w_est = 28, 28 # Defaults for CNN
-        num_pixels = 0
-
-        if preview_df is not None and isinstance(preview_df, pd.DataFrame) and not preview_df.empty:
-            cols = list(preview_df.columns)
-            if target_col in cols: cols.remove(target_col)
-            
-            if task == "Basic Image Classification":
-                num_pixels = len(cols)
-                if num_pixels > 0:
-                    dim_sqrt = int(math.sqrt(num_pixels))
-                    if dim_sqrt * dim_sqrt == num_pixels: img_h_est, img_w_est = dim_sqrt, dim_sqrt
-            else: # Tabular
-                num_cols = len([c for c in cols if pd.api.types.is_numeric_dtype(preview_df[c])])
-                cat_cols = [c for c in cols if pd.api.types.is_object_dtype(preview_df[c])]
-                one_hot_est = sum(min(10, preview_df[c].nunique(dropna=False)) for c in cat_cols)
-                input_dim_est = max(1, num_cols + one_hot_est)
-            
-            if target_col and target_col in preview_df.columns:
-                if task.endswith("Classification"):
-                    output_dim_est = max(1, preview_df[target_col].nunique(dropna=False))
-                    if output_dim_est == 2: output_dim_est = 1
-        else: logs += "Dataset preview not available for NN dimension estimation. Using defaults.\n"
-        return input_dim_est, output_dim_est, img_h_est, img_w_est, logs
-
-    def mlp_suggest_proxy_wrapper(target_range_str, current_logs, preview_df, target_col, task_type):
-        input_dim, output_dim, _, _, logs = get_data_dims_for_nn_suggestion(preview_df, target_col, task_type, current_logs)
-        logs += f"Using estimated input_dim: {input_dim}, output_dim: {output_dim} for MLP suggestion.\n"
-        suggested_str, logs = suggest_mlp_layers_for_range(input_dim, output_dim, target_range_str, logs)
-        param_count_str = "Error"
-        if suggested_str: param_count_str, logs = estimate_current_mlp_params(str(input_dim), suggested_str, str(output_dim), task_type, logs)
-        return suggested_str, logs, param_count_str
-    
-    pt_mlp_suggest_btn.click(fn=mlp_suggest_proxy_wrapper,
-        inputs=[pytorch_param_range_dd, current_logs_state, dataset_preview_df, target_column_name_txt, task_type_dd],
-        outputs=[pt_mlp_hidden_layers_txt, training_log_txt, pt_mlp_param_count_txt])
-
-    def mlp_estimate_proxy_wrapper(hidden_layers, current_logs, preview_df, target_col, task_type):
-        input_dim, output_dim, _, _, logs = get_data_dims_for_nn_suggestion(preview_df, target_col, task_type, current_logs)
-        logs += f"Using estimated input_dim: {input_dim}, output_dim: {output_dim} for MLP param estimation.\n"
-        param_count_str, logs = estimate_current_mlp_params(str(input_dim), hidden_layers, str(output_dim), task_type, logs)
-        return logs, param_count_str
-    
-    pt_mlp_estimate_params_btn.click(fn=mlp_estimate_proxy_wrapper,
-        inputs=[pt_mlp_hidden_layers_txt, current_logs_state, dataset_preview_df, target_column_name_txt, task_type_dd],
-        outputs=[training_log_txt, pt_mlp_param_count_txt])
-
-    def cnn_estimate_proxy_wrapper(current_logs, preview_df, target_col, task_type):
-        _, output_dim, img_h, img_w, logs = get_data_dims_for_nn_suggestion(preview_df, target_col, task_type, current_logs)
-        logs += f"Using estimated img_h: {img_h}, img_w: {img_w}, output_dim: {output_dim} for CNN param estimation.\n"
-        cnn_task_type = "classification" if task_type == "Basic Image Classification" else "regression"
-        param_count_str, logs = estimate_cnn_params(str(img_h), str(img_w), str(output_dim), cnn_task_type, logs)
-        return logs, param_count_str
-
-    pt_cnn_estimate_params_btn.click(fn=cnn_estimate_proxy_wrapper,
-        inputs=[current_logs_state, dataset_preview_df, target_column_name_txt, task_type_dd],
-        outputs=[training_log_txt, pt_cnn_param_count_txt])
-
-    def toggle_dataset_source_groups(source_choice):
-        return gr.update(visible=(source_choice == "Generate new dataset")), gr.update(visible=(source_choice == "Upload my own dataset (CSV, JSON, Parquet)"))
-    dataset_source_rb.change(fn=toggle_dataset_source_groups, inputs=dataset_source_rb, outputs=[generate_dataset_group, upload_dataset_group])
+    # Update output formats when family changes
     model_family_dd.change(fn=update_model_output_formats, inputs=model_family_dd, outputs=model_output_format_dd)
 
+    # Dataset generation button
     generate_dataset_btn.click(
         fn=generate_dataset_backend,
-        inputs=[task_type_dd, ds_gen_samples_num, ds_gen_features_num, ds_gen_classes_informative_num, 
-                ds_gen_format_dd, ds_gen_ai_suggest_cb, pytorch_param_range_dd, model_specific_dd, current_logs_state],
-        outputs=[dataset_preview_df, generated_data_state, training_log_txt, generated_dataset_download_file])
-
-    def process_uploaded_file(file_obj, logs_in):
-        logs, df_preview, stored_data_path = logs_in, None, None
-        if file_obj is None: logs += "Please upload a file first.\n"; return df_preview, logs, stored_data_path
-        logs += f"Uploaded file: {file_obj.name}\n"; stored_data_path = file_obj.name
-        try:
-            if file_obj.name.endswith(".csv"): df_preview = pd.read_csv(file_obj.name, nrows=5)
-            elif file_obj.name.endswith(".json"): df_preview = pd.read_json(file_obj.name, lines=True, nrows=5)
-            elif file_obj.name.endswith(".parquet"): temp_df = pd.read_parquet(file_obj.name); df_preview = temp_df.head()
-            logs += "Preview generated for uploaded file.\n"
-        except Exception as e: logs += f"Error previewing {file_obj.name}: {e}\n"
-        return df_preview, logs, stored_data_path
-    ds_upload_file.upload(fn=process_uploaded_file, inputs=[ds_upload_file, current_logs_state],
-                          outputs=[dataset_preview_df, training_log_txt, generated_data_state])
+        inputs=[task_type_dd, ds_gen_samples_num, ds_gen_features_num, ds_gen_classes_num, ds_gen_format_dd],
+        outputs=[dataset_preview_df, generated_data_state, training_log_txt, generated_dataset_download_file]
+    )
 
+    # Main training button
     train_model_btn.click(
         fn=train_model_wrapper,
-        inputs=[generated_data_state, target_column_name_txt, task_type_dd, model_family_dd, model_specific_dd,
-                pt_mlp_hidden_layers_txt, pt_mlp_activation_dd,
-                train_epochs_num, train_batch_size_num, train_learning_rate_num,
-                model_output_format_dd, training_log_txt],
-        outputs=[training_log_txt, evaluation_metrics_txt, download_trained_model_file, 
-                 model_param_count_output_txt, loss_plot_img])
+        inputs=[generated_data_state, target_column_name_txt, task_type_dd, model_family_dd, model_specific_dd, model_output_format_dd, training_log_txt],
+        outputs=[training_log_txt, evaluation_metrics_txt, download_trained_model_file, loss_plot_img]
+    )
 
+# Launch the application
 demo.queue().launch(debug=True, show_error=True)