Update app.py

app.py

@@ -37,6 +37,10 @@ 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')
+import matplotlib.pyplot as plt
 
 
 # --- Global Variables / Constants ---
@@ -93,13 +97,14 @@ class SimpleMLP(nn.Module):
         
         layers.append(nn.Linear(current_dim, output_dim))
         
-        if task_type == "classification": # Changed from task_type.endswith("Classification")
-            if output_dim == 1: layers.append(nn.Sigmoid()) # Binary
-            elif output_dim > 1: layers.append(nn.Softmax(dim=-1)) # Multi-class
+        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): # Added task_type to constructor for clarity
+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,
@@ -107,7 +112,7 @@ class SimpleCNN(nn.Module): # Added task_type to constructor for clarity
         super(SimpleCNN, self).__init__()
         self.input_channels = input_channels
         self.img_h, self.img_w = img_size_wh
-        self.num_classes = num_classes # This is the direct output dimension from the last linear layer
+        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()
@@ -115,9 +120,8 @@ class SimpleCNN(nn.Module): # Added task_type to constructor for clarity
 
         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 if pool1_k > 0 else h_out_conv1 # handle no pooling
-        w_pool1 = (w_out_conv1 - pool1_k) // pool1_s + 1 if pool1_k > 0 else w_out_conv1
-
+        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()
@@ -125,24 +129,23 @@ class SimpleCNN(nn.Module): # Added task_type to constructor for clarity
 
         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 if pool2_k > 0 else h_out_conv2
-        w_pool2 = (w_out_conv2 - pool2_k) // pool2_s + 1 if pool2_k > 0 else w_out_conv2
+        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. Current (h_pool2, w_pool2): ({h_pool2},{w_pool2}) from img ({self.img_h},{self.img_w})")
+            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) # Output layer before final activation
+        self.fc2 = nn.Linear(fc_hidden, self.num_classes)
         
+        # The final activation is now a separate attribute for clarity.
         if task_type == "classification":
-            if self.num_classes == 1: # Binary classification
+            if self.num_classes == 1: # Binary classification with BCELoss
                 self.final_activation = nn.Sigmoid()
-            elif self.num_classes > 1: # Multi-class classification
-                self.final_activation = nn.Softmax(dim=1)
-            else: # Should not happen for classification if num_classes is properly set
-                self.final_activation = nn.Identity()
+            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()
 
@@ -155,7 +158,7 @@ class SimpleCNN(nn.Module): # Added task_type to constructor for clarity
         x = self.final_activation(x)
         return x
 
-# --- Parameter Target Helpers (PARAM_RANGES, suggest_mlp_layers_for_range, estimate_current_mlp_params, estimate_cnn_params) ---
+# --- Parameter Target Helpers ---
 PARAM_RANGES = collections.OrderedDict([
     ("Tiny (<10k)", (0, 10000)),
     ("Small (10k-50k)", (10000, 50000)),
@@ -194,34 +197,32 @@ def suggest_mlp_layers_for_range(input_dim, output_dim, target_range_str, curren
     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=""): # Added task_type
+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
         
-        # Determine task_type for MLP constructor
         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=""): # Added task_type
+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
         
-        # Determine task_type for CNN constructor
-        cnn_task_type = "classification" if task_type.endswith("Classification") else "regression" # Assuming CNN for image is classification
+        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 (generate_dataset_backend, preprocess_tabular_data) ---
+# --- 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,
@@ -253,7 +254,7 @@ def generate_dataset_backend(task_type, n_samples_str, n_features_str,
     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)) # make_classification expects n_informative <= n_features
+            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)
@@ -280,15 +281,11 @@ def generate_dataset_backend(task_type, n_samples_str, n_features_str,
             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"
-            # For consistency, data_obj returned should be what train functions expect
-            # Sklearn train func can take df or path. PyTorch train func can take df, path, or (X,y) tuple.
-            # Returning df here for generated data is fine.
             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=""):
@@ -313,27 +310,6 @@ def preprocess_tabular_data(df_or_X, y_if_X_is_numpy, target_column_name, task_t
     ], remainder='passthrough')
 
     X_processed_np = preprocessor.fit_transform(X_df)
-    
-    feature_names_out_list = []
-    try: feature_names_out_list = list(preprocessor.get_feature_names_out())
-    except AttributeError: 
-        current_pos = 0
-        if numerical_features: feature_names_out_list.extend(numerical_features); current_pos += len(numerical_features)
-        if categorical_features:
-            cat_encoder = preprocessor.named_transformers_['cat'].named_steps['onehot']
-            if hasattr(cat_encoder, 'get_feature_names_out'):
-                cat_feature_names = cat_encoder.get_feature_names_out(categorical_features)
-            elif hasattr(cat_encoder, 'get_feature_names'):
-                 cat_feature_names = cat_encoder.get_feature_names(categorical_features)
-            else: # Estimate number of one-hot features
-                num_onehot_cols = X_processed_np.shape[1] - len(numerical_features) # Assuming only num and cat
-                cat_feature_names = [f"cat_feat_{i}" for i in range(num_onehot_cols)]
-            feature_names_out_list.extend(cat_feature_names); current_pos += len(cat_feature_names)
-        # Handle remainder='passthrough' if necessary, X_processed_np.shape[1] would be total
-        if X_processed_np.shape[1] > current_pos:
-            feature_names_out_list.extend([f"other_feat_{i}" for i in range(X_processed_np.shape[1] - current_pos)])
-
-
     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"
 
@@ -342,16 +318,16 @@ def preprocess_tabular_data(df_or_X, y_if_X_is_numpy, target_column_name, task_t
         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 # For NN output layer
+        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
     
-    return X_processed_np, y_processed_np, preprocessor, logs, processed_input_dim, output_dim_nn, feature_names_out_list
+    return X_processed_np, y_processed_np, preprocessor, logs, processed_input_dim, output_dim_nn
 
 
-# --- Training Functions (train_model_sklearn, train_model_pytorch) ---
+# --- 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"
@@ -367,11 +343,11 @@ def train_model_sklearn(data_input_obj, target_column, task_type, model_name, mo
     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: # check if target_column is empty
+    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, _, _, feature_names_original = preprocess_tabular_data(df, None, target_column, task_type, logs)
+        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)
@@ -414,21 +390,15 @@ def train_model_sklearn(data_input_obj, target_column, task_type, model_name, mo
             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_idx, col_name in enumerate(raw_X_for_types_df.columns):
+            for col_name in raw_X_for_types_df.columns:
                 col_dtype = raw_X_for_types_df[col_name].dtype
-                # Forcing float32 for numeric inputs to ONNX for broader compatibility
-                # ONNX is stricter about types than scikit-learn sometimes.
                 if pd.api.types.is_numeric_dtype(col_dtype):
-                    # Create a sample of the correct type for skl2onnx to infer shape and type
-                    # Shape [None, 1] implies one feature at a time for this column.
-                    # If a feature is multi-dimensional (e.g. embeddings), this needs adjustment.
-                    # For typical tabular, each column is one feature.
                     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])))
                 else:
-                    logs += f"Warning: Unsupported dtype {col_dtype} for column {col_name} in ONNX. Defaulting to FloatTensorType.\n"
-                    onnx_initial_types.append((col_name, FloatTensorType([None, 1]))) # Fallback
+                    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])))
             
             if not onnx_initial_types: raise ValueError("ONNX initial types failed: No valid columns found.")
             try:
@@ -437,9 +407,9 @@ def train_model_sklearn(data_input_obj, target_column, task_type, model_name, mo
                 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 with ONNX Runtime. Inputs: {[inp.name for inp in sess.get_inputs()]}\n"
+                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: # Fallback to PKL
+        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)
@@ -462,8 +432,7 @@ def train_model_pytorch(data_input_obj, target_column, task_type, model_type_pt,
             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) and len(data_input_obj) == 2 and isinstance(data_input_obj[0], np.ndarray) and isinstance(data_input_obj[1], np.ndarray):
-        X_numpy_for_pytorch, y_numpy_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:
@@ -475,36 +444,21 @@ def train_model_pytorch(data_input_obj, target_column, task_type, model_type_pt,
     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: # Check if target column is provided for DataFrame
-            logs += "Target column needed for MLP with DataFrame input.\n"; return logs, "MLP Target Error", None, "N/A", None
-
+        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
-            # target_column is only relevant if data_arg1 is a DataFrame
-            current_target_col = target_column if df_for_pytorch is not None else "target" # Placeholder if from numpy
-            
-            X_processed_np, y_processed_np, preprocessor_pipeline, logs, processed_input_dim_actual, nn_output_dim_actual, _ = \
+            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)":
-        if task_type != "Basic Image Classification": logs += "Warning: CNN selected, but task is not Basic Image Classification.\n"
-        
-        X_raw, y_raw = None, None
-        if df_for_pytorch is not None:
-            if not target_column or target_column not in df_for_pytorch.columns:
-                 logs += f"Target '{target_column}' not found/provided for CNN.\n"; return logs, "CNN Target Error", None, "N/A", None
-            X_raw = df_for_pytorch.drop(target_column, axis=1).values
-            y_raw = df_for_pytorch[target_column].values
-        elif X_numpy_for_pytorch is not None and y_numpy_for_pytorch is not None:
-            X_raw = X_numpy_for_pytorch; y_raw = y_numpy_for_pytorch
-        else: logs += "No valid data for CNN.\n"; return logs, "CNN Data Error", None, "N/A", None
-
+        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)
-        nn_output_dim_actual = len(le.classes_)
-        if nn_output_dim_actual == 2: nn_output_dim_actual = 1 
+        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))
@@ -513,86 +467,75 @@ def train_model_pytorch(data_input_obj, target_column, task_type, model_type_pt,
         
         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}, NN Output Dim: {nn_output_dim_actual}\n"
+        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
 
-    y_dtype = torch.float32 if (nn_output_dim_actual == 1 and task_type.endswith("Regression")) or \
-                              (nn_output_dim_actual == 1 and task_type.endswith("Classification")) \
-                           else torch.long
-    X_tensor = torch.tensor(X_processed_np, dtype=torch.float32)
-    y_tensor = torch.tensor(y_processed_np, dtype=y_dtype)
-    if nn_output_dim_actual == 1 and task_type.endswith("Classification"): y_tensor = y_tensor.unsqueeze(1)
-    if task_type.endswith("Regression"): y_tensor = y_tensor.unsqueeze(1)
+    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)
 
-    dataset = TensorDataset(X_tensor, y_tensor)
-    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0) 
+    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 task_type.endswith("Classification") else "regression")
+                                    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 task_type.endswith("Classification") else "regression") # Pass task_type
+                                      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"
-    if model_params_val > 500000: logs += "Warning: >500k params on CPU will be SLOW.\n"
 
-    is_classification_task = task_type.endswith("Classification") # Simplified condition
-    if is_classification_task:
-        criterion = nn.BCELoss() if nn_output_dim_actual == 1 else nn.CrossEntropyLoss()
-    else: criterion = nn.MSELoss()
+    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; num_batches = 0
-        for i, (batch_X, batch_y) in enumerate(dataloader): # Added enumerate for batch index
+        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(); num_batches += 1
-        avg_epoch_loss = epoch_loss_sum / num_batches if num_batches > 0 else 0
+            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)
-        logs += f"Epoch {epoch+1}/{epochs}, Avg Loss: {avg_epoch_loss:.4f}\n"
+        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():
-        if is_classification_task and dataloader.dataset and len(dataloader.dataset)>0 :
-            try:
-                all_preds, all_targets = [], []
-                for batch_X, batch_y in dataloader: # Evaluate on whole dataset (or a test split ideally)
-                    outputs = pytorch_model(batch_X)
-                    if nn_output_dim_actual == 1: predicted = (outputs > 0.5).float()
-                    else: _, predicted = torch.max(outputs.data, 1)
-                    all_preds.extend(predicted.cpu().numpy())
-                    all_targets.extend(batch_y.cpu().numpy())
-                
-                if all_targets and all_preds: # Check if lists are not empty
-                    # Ensure all_targets is 1D for accuracy_score if predicted is also 1D (binary case)
-                    all_targets_np = np.array(all_targets).squeeze()
-                    all_preds_np = np.array(all_preds).squeeze()
-                    acc = accuracy_score(all_targets_np, all_preds_np)
-                    metrics_out = f"Final Training Loss: {avg_epoch_loss:.4f}\nAccuracy on training data: {acc*100:.2f}%"
-                else:
-                    metrics_out = f"Final Training Loss: {avg_epoch_loss:.4f}\n (Could not compute accuracy)"
-
-            except Exception as eval_e: metrics_out = f"Final Training Loss: {avg_epoch_loss:.4f}\n Eval Error: {eval_e}"
-        else: metrics_out = f"Final Training Loss (MSE): {avg_epoch_loss:.4f}"
+        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:
-        import matplotlib # Use Agg backend for non-interactive environments like Spaces
-        matplotlib.use('Agg')
-        import matplotlib.pyplot as plt
         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"
@@ -602,23 +545,63 @@ def train_model_pytorch(data_input_obj, target_column, task_type, model_type_pt,
         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,
-            })
-            logs += f"PyTorch MLP (model + preprocessor) saved to {model_path_out}\n"
+            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}) # Save CNN architecture details
-            logs += f"PyTorch CNN (model state_dict + arch_details) saved to {model_path_out}\n"
-        else: logs += f"PyTorch {model_type_pt} (model state_dict) saved to {model_path_out}\n"
+            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}'. Saving as .pt\n"
-        model_path_out = get_temp_filepath(model_filename_base, "pt")
-        torch.save(pytorch_model.state_dict(), model_path_out)
+        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
+
+    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()}"
+        logs += error_msg + "\n"
+        return logs, error_msg, None, "N/A", None
+
+
 # --- Gradio UI Definition ---
 TASK_CHOICES = ["Tabular Classification", "Tabular Regression", "Basic Image Classification"]
 MODEL_FAMILIES = ["Scikit-learn (Classical ML)", "PyTorch (Neural Networks)"]
@@ -631,20 +614,19 @@ MODEL_OUTPUT_FORMATS_PYTORCH = [".pt (PyTorch)"]
 MLP_ACTIVATIONS = ["relu", "tanh", "sigmoid"]
 CLONE_GUIDE_TEXT = """
 ## How to Clone & Upgrade This Space for More Power:
-(Instructions as provided in previous response - omitted here for brevity but should be included)
+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!
 """
 
 # 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]
-initial_model_choices_for_specific_dd = []
 if _initial_family_default == "Scikit-learn (Classical ML)":
-    if _initial_task_default == "Tabular Classification": initial_model_choices_for_specific_dd = SKLEARN_MODELS_CLASSIFICATION
-    elif _initial_task_default == "Tabular Regression": initial_model_choices_for_specific_dd = SKLEARN_MODELS_REGRESSION
+    initial_model_choices_for_specific_dd = SKLEARN_MODELS_CLASSIFICATION
 elif _initial_family_default == "PyTorch (Neural Networks)":
-    if _initial_task_default.startswith("Tabular"): initial_model_choices_for_specific_dd = [PYTORCH_MODELS[0]]
-    elif _initial_task_default == "Basic Image Classification": initial_model_choices_for_specific_dd = [PYTORCH_MODELS[1]]
-initial_model_value_for_specific_dd = initial_model_choices_for_specific_dd[0] if initial_model_choices_for_specific_dd else None
+    initial_model_choices_for_specific_dd = [PYTORCH_MODELS[0]]
+initial_model_value_for_specific_dd = initial_model_choices_for_specific_dd[0]
 
 def update_model_options(task_choice, model_family_choice):
     choices, value = [], None
@@ -657,25 +639,18 @@ def update_model_options(task_choice, model_family_choice):
     value = choices[0] if choices else None
     return gr.update(choices=choices, value=value, visible=bool(choices))
 
-def update_param_range_visibility(model_family_choice):
-    return gr.update(visible=(model_family_choice == "PyTorch (Neural Networks)"))
-
-def update_pytorch_specific_options_visibility(model_choice_pytorch_family, specific_pytorch_model):
-    # Only proceed if family is PyTorch
-    if model_choice_pytorch_family != "PyTorch (Neural Networks)":
-        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) # Hide all: param_range, mlp_group, cnn_group
-    
-    param_range_visible = True # Always true if PyTorch family
-    is_mlp = (specific_pytorch_model == "Simple Neural Network (MLP)")
-    is_cnn = (specific_pytorch_model == "Simple Convolutional Network (CNN)")
-    return gr.update(visible=param_range_visible), gr.update(visible=is_mlp), gr.update(visible=is_cnn)
+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])
-    elif model_family_choice == "PyTorch (Neural Networks)": return gr.update(choices=MODEL_OUTPUT_FORMATS_PYTORCH, value=MODEL_OUTPUT_FORMATS_PYTORCH[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}""" # Hide footer too
+css = """.gradio-container { font-family: 'IBM Plex Sans', sans-serif; } footer {display:none !important}"""
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"), css=css) as demo:
     gr.Markdown("# 🧠 TrainAI ⚙️")
@@ -690,27 +665,28 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"),
                 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)
             
-            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], visible=(_initial_family_default == "PyTorch (Neural Networks)"))
-            with gr.Group(visible=(_initial_family_default == "PyTorch (Neural Networks)" and 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_family_default == "PyTorch (Neural Networks)" and 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)
-
-        with gr.TabItem("2. Configure Dataset"): # This tab should show generate_dataset_group by default
+            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)
+
+        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: # Default visible=True
+            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)
@@ -722,7 +698,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"),
             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'")
+            target_column_name_txt = gr.Textbox(label="Target Column Name (Case-Sensitive!)", placeholder="e.g., 'target' or 'label'", value="target")
             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)
 
@@ -731,7 +707,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"),
             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, step=1e-4, precision=6)
+                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])
             train_model_btn = gr.Button("🚀 Train Model", variant="primary")
             gr.Markdown("---")
@@ -743,46 +719,40 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"),
             download_trained_model_file = gr.File(label="Download Trained Model", interactive=False)
 
         with gr.TabItem("ℹ️ Guide & Info"):
-            # ... Guide content ... (omitted for brevity)
             gr.Markdown(CLONE_GUIDE_TEXT)
             
     # --- Event Handlers ---
     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=[pytorch_param_range_dd, 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=[pytorch_param_range_dd, pt_mlp_specific_group, pt_cnn_specific_group])
+    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:
-            temp_X_cols = [col for col in preview_df.columns if col != target_col] # Features
-            if not temp_X_cols and task == "Basic Image Classification": # Image data often has no named feature cols in preview
-                if preview_df.shape[1] == 1 and target_col in preview_df.columns: # Only target col
-                    pass # Cannot estimate image from only target
-                elif target_col in preview_df.columns:
-                     num_pixels = preview_df.shape[1] -1
-                else: # no target col (e.g. raw image data)
-                     num_pixels = preview_df.shape[1]
-
+            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: logs += f"Non-square image ({num_pixels} pixels) from preview. Using default {img_h_est}x{img_w_est} for suggestion.\n"
-                input_dim_est = img_h_est * img_w_est # For CNN, this is not input_dim to MLP but for SimpleCNN internal calcs
-            elif temp_X_cols: # Tabular
-                num_cols = len([col for col in temp_X_cols if pd.api.types.is_numeric_dtype(preview_df[col])])
-                cat_cols = [col for col in temp_X_cols if pd.api.types.is_object_dtype(preview_df[col])]
-                one_hot_est = sum(min(10, preview_df[col].nunique(dropna=False)) for col in cat_cols)
+            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 
+                    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
 
@@ -811,9 +781,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"),
     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"
-        # For CNN, task type for constructor is 'classification' typically
-        cnn_task_type_for_constructor = "classification" if task_type == "Basic Image Classification" else "regression" # Placeholder
-        param_count_str, logs = estimate_cnn_params(str(img_h), str(img_w), str(output_dim), cnn_task_type_for_constructor, logs)
+        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,
@@ -826,33 +795,32 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"),
     model_family_dd.change(fn=update_model_output_formats, inputs=model_family_dd, outputs=model_output_format_dd)
 
     generate_dataset_btn.click(
-        fn=generate_dataset_backend, # Pass the function directly
+        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, status_msg, stored_data_path = logs_in, None, "Upload failed or no file.", None
-        if file_obj is None: logs += "Please upload a file first.\n"; return df_preview, logs, status_msg, stored_data_path
+        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() # Parquet preview needs full read
-            status_msg = "Preview generated for uploaded file." if df_preview is not None else "Could not generate preview, but file is noted."
-            logs += status_msg + "\n"
-        except Exception as e: logs += f"Error previewing {file_obj.name}: {e}\n"; status_msg = f"Error previewing: {e}"
-        return df_preview, logs, status_msg, stored_data_path
+            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, training_log_txt, generated_data_state])
+                          outputs=[dataset_preview_df, training_log_txt, generated_data_state])
 
     train_model_btn.click(
-        fn=train_model_wrapper, # Wrapper function defined before UI
+        fn=train_model_wrapper,
         inputs=[generated_data_state, target_column_name_txt, task_type_dd, model_family_dd, model_specific_dd,
-                model_specific_dd, pt_mlp_hidden_layers_txt, pt_mlp_activation_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, download_trained_model_file])
+                 model_param_count_output_txt, loss_plot_img])
 
 demo.queue().launch(debug=True, show_error=True)