Update app.py

app.py

@@ -36,6 +36,10 @@ def setup_training():
     # Load your dataset with the specified configuration
     logging.info("Loading the MarbleX dataset")
     dataset = load_dataset("Oranblock/marblex_dataset", "config1")  # Replace "config1" with the appropriate config name
+
+    # Print dataset structure
+    logging.info(f"Dataset columns: {dataset['train'].column_names}")
+
     logging.info(f"Dataset loaded. Train size: {len(dataset['train'])}, Test size: {len(dataset['test'])}")
 
     # Load tokenizer and model