adds single token logic read/write , adds gpt-oss demo space , adds spaces refactor , adds new version of track tonic , adds logic in launch.sh

launch.sh

@@ -452,8 +452,10 @@ print_step "Step 1: User Authentication"
 echo "================================"
 
 print_info "You'll need two Hugging Face tokens:"
-echo "1. Write Token - Used during training for creating repositories and pushing models"
-echo "2. Read Token - Used in Trackio Space after training for security"
+echo "1. Write Token - Used initially for training and creating repositories"
+echo "2. Read Token - Will replace the write token in Trackio Space after training for security"
+echo ""
+print_info "The pipeline will start with the write token in HF_TOKEN, then switch to read token automatically."
 echo ""
 
 print_info "Getting Write Token (for training operations)..."
@@ -489,7 +491,7 @@ else
     exit 1
 fi
 
-# Set the main HF_TOKEN to write token for training operations
+# Set the main HF_TOKEN to write token for training operations (will be switched later)
 HF_TOKEN="$HF_WRITE_TOKEN"
 
 # Step 2: Select training configuration
@@ -669,8 +671,6 @@ fi
 
 # Set environment variables before creating virtual environment
 print_info "Setting up environment variables..."
-export HF_WRITE_TOKEN="$HF_WRITE_TOKEN"
-export HF_READ_TOKEN="$HF_READ_TOKEN"
 export HF_TOKEN="$HF_TOKEN"
 export TRACKIO_DATASET_REPO="$TRACKIO_DATASET_REPO"
 export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
@@ -682,8 +682,6 @@ source smollm3_env/bin/activate
 
 # Re-export environment variables in the virtual environment
 print_info "Configuring environment variables in virtual environment..."
-export HF_WRITE_TOKEN="$HF_WRITE_TOKEN"
-export HF_READ_TOKEN="$HF_READ_TOKEN"
 export HF_TOKEN="$HF_TOKEN"
 export TRACKIO_DATASET_REPO="$TRACKIO_DATASET_REPO"
 export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
@@ -712,16 +710,16 @@ print_status "HF token configured for Python API usage"
 print_info "Username: $HF_USERNAME (auto-detected from token)"
 print_info "Token available in environment: ${HF_TOKEN:0:10}...${HF_TOKEN: -4}"
 
-# Verify tokens are available in the virtual environment
+# Verify token is available in the virtual environment
 print_info "Verifying token availability in virtual environment..."
-if [ -n "$HF_WRITE_TOKEN" ] && [ -n "$HF_READ_TOKEN" ] && [ -n "$HUGGING_FACE_HUB_TOKEN" ]; then
-    print_status "✅ Tokens properly configured in virtual environment"
-    print_info "  HF_WRITE_TOKEN: ${HF_WRITE_TOKEN:0:10}...${HF_WRITE_TOKEN: -4}"
-    print_info "  HF_READ_TOKEN: ${HF_READ_TOKEN:0:10}...${HF_READ_TOKEN: -4}"
+if [ -n "$HF_TOKEN" ] && [ -n "$HUGGING_FACE_HUB_TOKEN" ]; then
+    print_status "✅ Token properly configured in virtual environment"
+    print_info "  HF_TOKEN: ${HF_TOKEN:0:10}...${HF_TOKEN: -4} (currently using WRITE token)"
     print_info "  HUGGING_FACE_HUB_TOKEN: ${HUGGING_FACE_HUB_TOKEN:0:10}...${HUGGING_FACE_HUB_TOKEN: -4}"
+    print_info "  Will be switched to READ token after training for security"
 else
-    print_error "❌ Tokens not properly configured in virtual environment"
-    print_error "Please check your tokens and try again"
+    print_error "❌ Token not properly configured in virtual environment"
+    print_error "Please check your token and try again"
     exit 1
 fi
 
@@ -771,8 +769,6 @@ print_info "Username will be auto-detected from token"
 print_info "Secrets will be set automatically via API"
 
 # Ensure environment variables are available for the script
-export HF_WRITE_TOKEN="$HF_WRITE_TOKEN"
-export HF_READ_TOKEN="$HF_READ_TOKEN"
 export HF_TOKEN="$HF_TOKEN"
 export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
 export HF_USERNAME="$HF_USERNAME"
@@ -792,8 +788,6 @@ print_info "Username will be auto-detected from token"
 print_info "Dataset repository: $TRACKIO_DATASET_REPO"
 
 # Ensure environment variables are available for the script
-export HF_WRITE_TOKEN="$HF_WRITE_TOKEN"
-export HF_READ_TOKEN="$HF_READ_TOKEN"
 export HF_TOKEN="$HF_TOKEN"
 export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
 export HF_USERNAME="$HF_USERNAME"
@@ -809,8 +803,6 @@ print_info "Configuring Trackio ..."
 print_info "Username will be auto-detected from token"
 
 # Ensure environment variables are available for the script
-export HF_WRITE_TOKEN="$HF_WRITE_TOKEN"
-export HF_READ_TOKEN="$HF_READ_TOKEN"
 export HF_TOKEN="$HF_TOKEN"
 export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
 export HF_USERNAME="$HF_USERNAME"
@@ -920,7 +912,7 @@ fi
 print_step "Step 16.5: Switching to Read Token for Security"
 echo "===================================================="
 
-print_info "Switching Trackio Space from write token to read token for security..."
+print_info "Switching Trackio Space HF_TOKEN from write token to read token for security..."
 print_info "This ensures the space can only read datasets, not write to repositories"
 
 # Ensure environment variables are available for token switch
@@ -928,12 +920,12 @@ export HF_TOKEN="$HF_WRITE_TOKEN"  # Use write token to update space
 export HUGGING_FACE_HUB_TOKEN="$HF_WRITE_TOKEN"
 export HF_USERNAME="$HF_USERNAME"
 
-# Switch to read token in Trackio Space
+# Switch HF_TOKEN in Trackio Space from write to read token
 cd scripts/trackio_tonic
 python switch_to_read_token.py "$HF_USERNAME/$TRACKIO_SPACE_NAME" "$HF_READ_TOKEN" "$HF_WRITE_TOKEN"
 
 if [ $? -eq 0 ]; then
-    print_status "✅ Successfully switched Trackio Space to read token"
+    print_status "✅ Successfully switched Trackio Space HF_TOKEN to read token"
     print_info "🔒 Space now uses read-only permissions for security"
 else
     print_warning "⚠️ Failed to switch to read token, but continuing with pipeline"
@@ -957,8 +949,6 @@ if [ "$DEPLOY_DEMO" = "y" ] || [ "$DEPLOY_DEMO" = "Y" ]; then
     DEMO_SUBFOLDER=""
     
     # Ensure environment variables are available for demo deployment
-export HF_WRITE_TOKEN="$HF_WRITE_TOKEN"
-export HF_READ_TOKEN="$HF_READ_TOKEN"
 export HF_TOKEN="$HF_TOKEN"
 export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
 export HF_USERNAME="$HF_USERNAME"
@@ -999,7 +989,7 @@ cat > training_summary.md << EOF
 - **HF Dataset**: $TRACKIO_DATASET_REPO
 - **Training Config**: $TRAINING_CONFIG_TYPE
 - **Trainer Type**: $TRAINER_TYPE
-- **Security**: Dual token system (write + read tokens)
+- **Security**: Single HF_TOKEN switched from write to read token
 $(if [ "$TRAINING_CONFIG_TYPE" = "H100 Lightweight (Rapid)" ]; then
 echo "- **Dataset Sample Size**: ${DATASET_SAMPLE_SIZE:-80000}"
 fi)
@@ -1015,7 +1005,7 @@ fi)
 - **Model Repository**: https://huggingface.co/$REPO_NAME
 - **Trackio Monitoring**: $TRACKIO_URL
 - **Experiment Data**: https://huggingface.co/datasets/$TRACKIO_DATASET_REPO
-- **Security**: Trackio Space switched to read-only token for security
+- **Security**: Trackio Space HF_TOKEN switched to read-only token for security
 $(if [ "$DEPLOY_DEMO" = "y" ] || [ "$DEPLOY_DEMO" = "Y" ]; then
 echo "- **Demo Space**: https://huggingface.co/spaces/$HF_USERNAME/${REPO_NAME}-demo"
 fi)
@@ -1053,7 +1043,7 @@ echo ""
 echo "🚀 Next steps:"
 echo "1. Monitor training progress in your Trackio Space"
 echo "2. Check the model repository on Hugging Face Hub"
-echo "3. Your Trackio Space is now secured with read-only permissions"
+echo "3. Your Trackio Space HF_TOKEN is now secured with read-only permissions"
 $(if [ "$DEPLOY_DEMO" = "y" ] || [ "$DEPLOY_DEMO" = "Y" ]; then
 echo "3. Make your huggingface space a ZeroGPU Space & Test your model"
 fi)

scripts/dataset_tonic/setup_hf_dataset.py

@@ -145,7 +145,7 @@ def setup_trackio_dataset(dataset_name: str = None, token: str = None) -> bool:
 
 def add_initial_experiment_data(repo_id: str, token: str = None) -> bool:
     """
-    Add initial experiment data to the dataset.
+    Add initial experiment data to the dataset using data preservation.
     
     Args:
         repo_id (str): Dataset repository ID
@@ -163,89 +163,95 @@ def add_initial_experiment_data(repo_id: str, token: str = None) -> bool:
             print("⚠️  No token available for uploading data")
             return False
         
-        # Initial experiment data
-        initial_experiments = [
-            {
-                'experiment_id': f'exp_{datetime.now().strftime("%Y%m%d_%H%M%S")}',
-                'name': 'smollm3-finetune-demo',
-                'description': 'SmolLM3 fine-tuning experiment demo with comprehensive metrics tracking',
-                'created_at': datetime.now().isoformat(),
-                'status': 'completed',
-                'metrics': json.dumps([
-                    {
-                        'timestamp': datetime.now().isoformat(),
-                        'step': 100,
-                        'metrics': {
-                            'loss': 1.15,
-                            'grad_norm': 10.5,
-                            'learning_rate': 5e-6,
-                            'num_tokens': 1000000.0,
-                            'mean_token_accuracy': 0.76,
-                            'epoch': 0.1,
-                            'total_tokens': 1000000.0,
-                            'throughput': 2000000.0,
-                            'step_time': 0.5,
-                            'batch_size': 2,
-                            'seq_len': 4096,
-                            'token_acc': 0.76,
-                            'gpu_memory_allocated': 15.2,
-                            'gpu_memory_reserved': 70.1,
-                            'gpu_utilization': 85.2,
-                            'cpu_percent': 2.7,
-                            'memory_percent': 10.1
-                        }
-                    }
-                ]),
-                'parameters': json.dumps({
-                    'model_name': 'HuggingFaceTB/SmolLM3-3B',
-                    'max_seq_length': 4096,
-                    'batch_size': 2,
-                    'learning_rate': 5e-6,
-                    'epochs': 3,
-                    'dataset': 'OpenHermes-FR',
-                    'trainer_type': 'SFTTrainer',
-                    'hardware': 'GPU (H100/A100)',
-                    'mixed_precision': True,
-                    'gradient_checkpointing': True,
-                    'flash_attention': True
-                }),
-                'artifacts': json.dumps([]),
-                'logs': json.dumps([
-                    {
-                        'timestamp': datetime.now().isoformat(),
-                        'level': 'INFO',
-                        'message': 'Training started successfully'
-                    },
-                    {
-                        'timestamp': datetime.now().isoformat(),
-                        'level': 'INFO',
-                        'message': 'Model loaded and configured'
-                    },
-                    {
-                        'timestamp': datetime.now().isoformat(),
-                        'level': 'INFO',
-                        'message': 'Dataset loaded and preprocessed'
-                    }
-                ]),
-                'last_updated': datetime.now().isoformat()
-            }
-        ]
+        # Import dataset manager
+        import sys
+        sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
+        from dataset_utils import TrackioDatasetManager
         
-        # Create dataset and upload
-        from datasets import Dataset
+        # Initialize dataset manager
+        dataset_manager = TrackioDatasetManager(repo_id, token)
         
-        # Create dataset from the initial experiments
-        dataset = Dataset.from_list(initial_experiments)
+        # Check if dataset already has data
+        existing_experiments = dataset_manager.load_existing_experiments()
+        if existing_experiments:
+            print(f"ℹ️  Dataset already contains {len(existing_experiments)} experiments, preserving existing data")
         
-        # Push to hub
-        dataset.push_to_hub(
-            repo_id,
-            token=token,
-            private=False,
-            commit_message="Add initial experiment data"
-        )
+        # Initial experiment data
+        initial_experiment = {
+            'experiment_id': f'exp_demo_{datetime.now().strftime("%Y%m%d_%H%M%S")}',
+            'name': 'smollm3-finetune-demo',
+            'description': 'SmolLM3 fine-tuning experiment demo with comprehensive metrics tracking',
+            'created_at': datetime.now().isoformat(),
+            'status': 'completed',
+            'metrics': json.dumps([
+                {
+                    'timestamp': datetime.now().isoformat(),
+                    'step': 100,
+                    'metrics': {
+                        'loss': 1.15,
+                        'grad_norm': 10.5,
+                        'learning_rate': 5e-6,
+                        'num_tokens': 1000000.0,
+                        'mean_token_accuracy': 0.76,
+                        'epoch': 0.1,
+                        'total_tokens': 1000000.0,
+                        'throughput': 2000000.0,
+                        'step_time': 0.5,
+                        'batch_size': 2,
+                        'seq_len': 4096,
+                        'token_acc': 0.76,
+                        'gpu_memory_allocated': 15.2,
+                        'gpu_memory_reserved': 70.1,
+                        'gpu_utilization': 85.2,
+                        'cpu_percent': 2.7,
+                        'memory_percent': 10.1
+                    }
+                }
+            ]),
+            'parameters': json.dumps({
+                'model_name': 'HuggingFaceTB/SmolLM3-3B',
+                'max_seq_length': 4096,
+                'batch_size': 2,
+                'learning_rate': 5e-6,
+                'epochs': 3,
+                'dataset': 'OpenHermes-FR',
+                'trainer_type': 'SFTTrainer',
+                'hardware': 'GPU (H100/A100)',
+                'mixed_precision': True,
+                'gradient_checkpointing': True,
+                'flash_attention': True
+            }),
+            'artifacts': json.dumps([]),
+            'logs': json.dumps([
+                {
+                    'timestamp': datetime.now().isoformat(),
+                    'level': 'INFO',
+                    'message': 'Training started successfully'
+                },
+                {
+                    'timestamp': datetime.now().isoformat(),
+                    'level': 'INFO',
+                    'message': 'Model loaded and configured'
+                },
+                {
+                    'timestamp': datetime.now().isoformat(),
+                    'level': 'INFO',
+                    'message': 'Dataset loaded and preprocessed'
+                }
+            ]),
+            'last_updated': datetime.now().isoformat()
+        }
+        
+        # Use dataset manager to safely add the experiment
+        success = dataset_manager.upsert_experiment(initial_experiment)
         
-        print(f"✅ Successfully uploaded initial experiment data to {repo_id}")
+        if success:
+            print(f"✅ Successfully added initial experiment data to {repo_id}")
+            final_count = len(dataset_manager.load_existing_experiments())
+            print(f"📊 Dataset now contains {final_count} total experiments")
+        else:
+            print(f"❌ Failed to add initial experiment data to {repo_id}")
+            return False
         
         # Add README template
         add_dataset_readme(repo_id, token)

scripts/deploy_demo_space.py

@@ -38,7 +38,8 @@ class DemoSpaceDeployer:
     """Deploy demo space to Hugging Face Spaces"""
     
     def __init__(self, hf_token: str, hf_username: str, model_id: str, 
-                 subfolder: str = "int4", space_name: Optional[str] = None):
+                 subfolder: str = "int4", space_name: Optional[str] = None, 
+                 demo_type: Optional[str] = None):
         self.hf_token = hf_token
         self.hf_username = hf_username
         self.model_id = model_id
@@ -47,8 +48,13 @@ class DemoSpaceDeployer:
         self.space_id = f"{hf_username}/{self.space_name}"
         self.space_url = f"https://huggingface.co/spaces/{self.space_id}"
         
-        # Template paths
-        self.template_dir = Path(__file__).parent.parent / "templates" / "spaces" / "demo"
+        # Determine demo type from model_id if not provided
+        if demo_type is None:
+            demo_type = self._detect_demo_type(model_id)
+        
+        # Template paths based on model type
+        self.demo_type = demo_type
+        self.template_dir = Path(__file__).parent.parent / "templates" / "spaces" / f"demo_{demo_type}"
         self.workspace_dir = Path.cwd()
         
         # Initialize HF API
@@ -58,6 +64,107 @@ class DemoSpaceDeployer:
             self.api = None
             logger.warning("huggingface_hub not available, using CLI fallback")
     
+    def _detect_demo_type(self, model_id: str) -> str:
+        """Detect the appropriate demo type based on model ID"""
+        model_id_lower = model_id.lower()
+        
+        # Check for GPT-OSS models
+        if "gpt-oss" in model_id_lower or "gpt_oss" in model_id_lower:
+            logger.info(f"Detected GPT-OSS model, using demo_gpt template")
+            return "gpt"
+        
+        # Check for SmolLM models (default)
+        elif "smollm" in model_id_lower or "smol" in model_id_lower:
+            logger.info(f"Detected SmolLM model, using demo_smol template")
+            return "smol"
+        
+        # Default to SmolLM for unknown models
+        else:
+            logger.info(f"Unknown model type, defaulting to demo_smol template")
+            return "smol"
+    
+    def _generate_env_setup(self) -> str:
+        """Generate environment variable setup based on demo type and model"""
+        if self.demo_type == "gpt":
+            # For GPT-OSS models, we need more sophisticated environment setup
+            model_name = self.model_id.split("/")[-1] if "/" in self.model_id else self.model_id
+            
+            env_setup = f"""
+# Environment variables for GPT-OSS model configuration
+import os
+os.environ['HF_MODEL_ID'] = '{self.model_id}'
+os.environ['LORA_MODEL_ID'] = '{self.model_id}'
+os.environ['BASE_MODEL_ID'] = 'openai/gpt-oss-20b'
+os.environ['MODEL_SUBFOLDER'] = '{self.subfolder if self.subfolder else ""}'
+os.environ['MODEL_NAME'] = '{model_name}'
+
+"""
+        else:
+            # For SmolLM models, use simpler setup
+            env_setup = f"""
+# Environment variables for model configuration
+import os
+os.environ['HF_MODEL_ID'] = '{self.model_id}'
+os.environ['MODEL_SUBFOLDER'] = '{self.subfolder if self.subfolder else ""}'
+os.environ['MODEL_NAME'] = '{self.model_id.split("/")[-1]}'
+
+"""
+        return env_setup
+    
+    def _set_model_variables(self):
+        """Set model-specific environment variables in the space"""
+        try:
+            # Common variables for all models
+            self.api.add_space_variable(
+                repo_id=self.space_id,
+                key="HF_MODEL_ID",
+                value=self.model_id,
+                description="Model ID for the demo"
+            )
+            logger.info(f"✅ Successfully set HF_MODEL_ID variable: {self.model_id}")
+            
+            if self.subfolder and self.subfolder.strip():
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="MODEL_SUBFOLDER",
+                    value=self.subfolder,
+                    description="Model subfolder for the demo"
+                )
+                logger.info(f"✅ Successfully set MODEL_SUBFOLDER variable: {self.subfolder}")
+            else:
+                logger.info("ℹ️ No subfolder specified, using main model")
+            
+            # GPT-OSS specific variables
+            if self.demo_type == "gpt":
+                model_name = self.model_id.split("/")[-1] if "/" in self.model_id else self.model_id
+                
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="LORA_MODEL_ID",
+                    value=self.model_id,
+                    description="LoRA/Fine-tuned model ID"
+                )
+                logger.info(f"✅ Successfully set LORA_MODEL_ID variable: {self.model_id}")
+                
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="BASE_MODEL_ID",
+                    value="openai/gpt-oss-20b",
+                    description="Base model ID for GPT-OSS"
+                )
+                logger.info("✅ Successfully set BASE_MODEL_ID variable: openai/gpt-oss-20b")
+                
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="MODEL_NAME",
+                    value=model_name,
+                    description="Display name for the model"
+                )
+                logger.info(f"✅ Successfully set MODEL_NAME variable: {model_name}")
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to set model variables: {e}")
+    
     def validate_model_exists(self) -> bool:
         """Validate that the model exists on Hugging Face Hub"""
         try:
@@ -187,14 +294,7 @@ class DemoSpaceDeployer:
                     content = f.read()
                 
                 # Add environment variable setup at the top
-                env_setup = f"""
-# Environment variables for model configuration
-import os
-os.environ['HF_MODEL_ID'] = '{self.model_id}'
-os.environ['MODEL_SUBFOLDER'] = '{self.subfolder if self.subfolder else ""}'
-os.environ['MODEL_NAME'] = '{self.model_id.split("/")[-1]}'
-
-"""
+                env_setup = self._generate_env_setup()
                 
                 # Insert after imports
                 lines = content.split('\n')
@@ -335,24 +435,7 @@ Simply start chatting with the model using the interface below!
                 logger.info("✅ Successfully set HF_TOKEN secret via API")
                 
                 # Set model-specific environment variables
-                self.api.add_space_variable(
-                    repo_id=self.space_id,
-                    key="HF_MODEL_ID",
-                    value=self.model_id,
-                    description="Model ID for the demo"
-                )
-                logger.info(f"✅ Successfully set HF_MODEL_ID variable: {self.model_id}")
-                
-                if self.subfolder and self.subfolder.strip():
-                    self.api.add_space_variable(
-                        repo_id=self.space_id,
-                        key="MODEL_SUBFOLDER",
-                        value=self.subfolder,
-                        description="Model subfolder for the demo"
-                    )
-                    logger.info(f"✅ Successfully set MODEL_SUBFOLDER variable: {self.subfolder}")
-                else:
-                    logger.info("ℹ️ No subfolder specified, using main model")
+                self._set_model_variables()
                 
                 return True
                 
@@ -375,6 +458,13 @@ Simply start chatting with the model using the interface below!
         else:
             logger.info("   MODEL_SUBFOLDER=(empty - using main model)")
         
+        # GPT-OSS specific variables
+        if self.demo_type == "gpt":
+            model_name = self.model_id.split("/")[-1] if "/" in self.model_id else self.model_id
+            logger.info(f"   LORA_MODEL_ID={self.model_id}")
+            logger.info(f"   BASE_MODEL_ID=openai/gpt-oss-20b")
+            logger.info(f"   MODEL_NAME={model_name}")
+        
         logger.info(f"\n🔧 To set secrets in your Space:")
         logger.info(f"1. Go to your Space settings: {self.space_url}/settings")
         logger.info("2. Navigate to the 'Repository secrets' section")
@@ -389,6 +479,17 @@ Simply start chatting with the model using the interface below!
         else:
             logger.info("   Name: MODEL_SUBFOLDER")
             logger.info("   Value: (leave empty)")
+        
+        # GPT-OSS specific variables
+        if self.demo_type == "gpt":
+            model_name = self.model_id.split("/")[-1] if "/" in self.model_id else self.model_id
+            logger.info(f"   Name: LORA_MODEL_ID")
+            logger.info(f"   Value: {self.model_id}")
+            logger.info(f"   Name: BASE_MODEL_ID")
+            logger.info(f"   Value: openai/gpt-oss-20b")
+            logger.info(f"   Name: MODEL_NAME")
+            logger.info(f"   Value: {model_name}")
+        
         logger.info("4. Save the secrets")
         
         return True
@@ -471,6 +572,7 @@ def main():
     parser.add_argument("--model-id", required=True, help="Model ID to deploy demo for")
     parser.add_argument("--subfolder", default="int4", help="Model subfolder (default: int4)")
     parser.add_argument("--space-name", help="Custom space name (optional)")
+    parser.add_argument("--demo-type", choices=["smol", "gpt"], help="Demo type: 'smol' for SmolLM, 'gpt' for GPT-OSS (auto-detected if not specified)")
     
     args = parser.parse_args()
     
@@ -479,7 +581,8 @@ def main():
         hf_username=args.hf_username,
         model_id=args.model_id,
         subfolder=args.subfolder,
-        space_name=args.space_name
+        space_name=args.space_name,
+        demo_type=args.demo_type
     )
     
     success = deployer.deploy()

scripts/model_tonic/push_gpt_oss_to_huggingface.py

@@ -169,8 +169,8 @@ If you use this model in your research, please cite:
   author = {{{author_name}}},
   title = {{{model_name}}},
   year = {{{datetime.now().year}}},
-  publisher = {Hugging Face},
-  journal = {Hugging Face repository},
+  publisher = {{Hugging Face}},
+  journal = {{Hugging Face repository}},
   howpublished = {{\\url{{https://huggingface.co/{model_name}}}}}
   }}
 ```

scripts/trackio_tonic/configure_trackio.py

@@ -79,13 +79,11 @@ def configure_trackio():
     print("🔧 Trackio Configuration")
     print("=" * 40)
     
-    # Get HF tokens and user info
-    hf_write_token = os.environ.get('HF_WRITE_TOKEN')
-    hf_read_token = os.environ.get('HF_READ_TOKEN')
-    hf_token = os.environ.get('HF_TOKEN')  # Legacy support
+    # Get HF token (single token approach)
+    hf_token = os.environ.get('HF_TOKEN')
     
-    # Use write token if available, otherwise fall back to HF_TOKEN
-    active_token = hf_write_token or hf_token
+    # Use the single HF_TOKEN
+    active_token = hf_token
     
     if active_token:
         username = get_username_from_token(active_token)
@@ -102,9 +100,7 @@ def configure_trackio():
     
     # Current configuration
     current_config = {
-        'HF_WRITE_TOKEN': hf_write_token or 'Not set',
-        'HF_READ_TOKEN': hf_read_token or 'Not set',
-        'HF_TOKEN': hf_token or 'Not set',  # Legacy
+        'HF_TOKEN': hf_token or 'Not set',
         'TRACKIO_DATASET_REPO': dataset_repo,
         'SPACE_ID': os.environ.get('SPACE_ID', 'Not set'),
         'TRACKIO_URL': os.environ.get('TRACKIO_URL', 'Not set')
@@ -116,12 +112,10 @@ def configure_trackio():
         print(f"   {status} {key}: {value}")
     
     print("\n🎯 Configuration Options:")
-    print("1. Set HF_WRITE_TOKEN - Required for training operations")
-    print("2. Set HF_READ_TOKEN - Required for Trackio Space security")
-    print("3. Set HF_TOKEN - Legacy token (fallback)")
-    print("4. Set TRACKIO_DATASET_REPO - Dataset repository (optional)")
-    print("5. Set SPACE_ID - HF Space ID (auto-detected)")
-    print("6. Set TRACKIO_URL - Trackio Space URL (auto-detected)")
+    print("1. Set HF_TOKEN - Main token (starts as write, switches to read after training)")
+    print("2. Set TRACKIO_DATASET_REPO - Dataset repository (optional)")
+    print("3. Set SPACE_ID - HF Space ID (auto-detected)")
+    print("4. Set TRACKIO_URL - Trackio Space URL (auto-detected)")
     
     # Check if running on HF Spaces
     if os.environ.get('SPACE_ID'):
@@ -131,37 +125,21 @@ def configure_trackio():
     # Validate configuration
     print("\n🔍 Configuration Validation:")
     
-    # Check HF_WRITE_TOKEN
-    if current_config['HF_WRITE_TOKEN'] != 'Not set':
-        print("✅ HF_WRITE_TOKEN is set")
-        print("   This allows training operations and repository creation")
-    else:
-        print("❌ HF_WRITE_TOKEN is not set")
-        print("   Please set HF_WRITE_TOKEN for training operations")
-        print("   Get your token from: https://huggingface.co/settings/tokens")
-    
-    # Check HF_READ_TOKEN
-    if current_config['HF_READ_TOKEN'] != 'Not set':
-        print("✅ HF_READ_TOKEN is set")
-        print("   This will be used for Trackio Space security")
-    else:
-        print("❌ HF_READ_TOKEN is not set")
-        print("   Please set HF_READ_TOKEN for Space security")
-        print("   Get your token from: https://huggingface.co/settings/tokens")
-    
-    # Check legacy HF_TOKEN
+    # Check HF_TOKEN
     if current_config['HF_TOKEN'] != 'Not set':
-        print("✅ HF_TOKEN (legacy) is set")
-        print("   This provides fallback functionality")
+        print("✅ HF_TOKEN is set")
+        print("   This allows training operations and dataset access")
+        print("   Note: Token will be automatically switched from write to read after training")
     else:
-        print("⚠️  HF_TOKEN (legacy) is not set")
-        print("   This is optional if using HF_WRITE_TOKEN")
+        print("❌ HF_TOKEN is not set")
+        print("   Please set HF_TOKEN for training operations")
+        print("   Get your token from: https://huggingface.co/settings/tokens")
     
     # Check dataset repository
     print(f"📊 Dataset Repository: {dataset_repo}")
     
     # Test dataset access if token is available
-    test_token = current_config['HF_WRITE_TOKEN'] or current_config['HF_TOKEN']
+    test_token = current_config['HF_TOKEN']
     if test_token != 'Not set':
         print("\n🧪 Testing Dataset Access...")
         try:
@@ -216,15 +194,13 @@ def configure_trackio():
     # Generate configuration file
     config_file = "trackio_config.json"
     config_data = {
-        'hf_write_token': current_config['HF_WRITE_TOKEN'],
-        'hf_read_token': current_config['HF_READ_TOKEN'],
-        'hf_token': current_config['HF_TOKEN'],  # Legacy
+        'hf_token': current_config['HF_TOKEN'],
         'dataset_repo': current_config['TRACKIO_DATASET_REPO'],
         'space_id': current_config['SPACE_ID'],
         'trackio_url': current_config['TRACKIO_URL'],
         'username': username,
         'last_updated': datetime.now().isoformat(),
-        'notes': 'Trackio configuration - set these as environment variables in your HF Space'
+        'notes': 'Trackio configuration - HF_TOKEN starts as write token, switches to read token after training'
     }
     
     with open(config_file, 'w') as f:
@@ -235,16 +211,14 @@ def configure_trackio():
     # Show environment variable commands
     print("\n📝 Environment Variables for HF Space:")
     print("=" * 50)
-    print(f"HF_WRITE_TOKEN={current_config['HF_WRITE_TOKEN']}")
-    print(f"HF_READ_TOKEN={current_config['HF_READ_TOKEN']}")
-    print(f"HF_TOKEN={current_config['HF_TOKEN']}")  # Legacy
+    print(f"HF_TOKEN={current_config['HF_TOKEN']}")
     print(f"TRACKIO_DATASET_REPO={current_config['TRACKIO_DATASET_REPO']}")
     if current_config['TRACKIO_URL'] != 'Not set':
         print(f"TRACKIO_URL={current_config['TRACKIO_URL']}")
     
     print("\n🎯 Next Steps:")
-    print("1. Set HF_WRITE_TOKEN in your HF Space environment variables")
-    print("2. Set HF_READ_TOKEN in your HF Space environment variables")
+    print("1. HF_TOKEN will be automatically set during deployment (starts as write token)")
+    print("2. HF_TOKEN will be automatically switched to read token after training")
     print("3. Optionally set TRACKIO_DATASET_REPO to use a different dataset")
     print("4. Deploy your updated app.py to the Space")
     print("5. Run setup_hf_dataset.py if you haven't created the dataset yet")

scripts/trackio_tonic/deploy_trackio_space.py

@@ -196,16 +196,16 @@ class TrackioSpaceDeployer:
             
             # Get the project root directory (3 levels up from this script)
             project_root = Path(__file__).parent.parent.parent
-            templates_dir = project_root / "templates" / "spaces"
+            templates_dir = project_root / "templates" / "spaces" / "trackio"
             
-            # Files to copy from templates/spaces
+            # Files to copy from templates/spaces/trackio
             files_to_copy = [
                 "app.py",
                 "requirements.txt",
                 "README.md"
             ]
             
-            # Copy files from templates/spaces to temp directory
+            # Copy files from templates/spaces/trackio to temp directory
             copied_files = []
             for file_name in files_to_copy:
                 source_path = templates_dir / file_name
@@ -334,36 +334,16 @@ class TrackioSpaceDeployer:
             
             repo_id = f"{self.username}/{self.space_name}"
             
-            # Get the HF tokens from environment or use the provided token
-            hf_write_token = os.getenv('HF_WRITE_TOKEN', self.token)
-            hf_read_token = os.getenv('HF_READ_TOKEN', self.token)
-            hf_token = os.getenv('HF_TOKEN', self.token)  # Legacy
+            # Use the provided token as HF_TOKEN (starts as write token, will be switched to read token later)
+            hf_token = self.token
             
-            # Set the HF_WRITE_TOKEN secret for the space using the API
+            # Set the HF_TOKEN secret for the space using the API
             try:
-                self.api.add_space_secret(
-                    repo_id=repo_id,
-                    key="HF_WRITE_TOKEN",
-                    value=hf_write_token,
-                    description="Hugging Face write token for training operations"
-                )
-                print("✅ Successfully set HF_WRITE_TOKEN secret via API")
-                
-                # Set the HF_READ_TOKEN secret for the space using the API
-                self.api.add_space_secret(
-                    repo_id=repo_id,
-                    key="HF_READ_TOKEN",
-                    value=hf_read_token,
-                    description="Hugging Face read token for security"
-                )
-                print("✅ Successfully set HF_READ_TOKEN secret via API")
-                
-                # Set legacy HF_TOKEN secret for backward compatibility
                 self.api.add_space_secret(
                     repo_id=repo_id,
                     key="HF_TOKEN",
                     value=hf_token,
-                    description="Hugging Face token for dataset access (legacy)"
+                    description="Hugging Face token for dataset access (starts as write, switches to read)"
                 )
                 print("✅ Successfully set HF_TOKEN secret via API")
                 
@@ -401,13 +381,9 @@ class TrackioSpaceDeployer:
         """Fallback method for manual secret setup"""
         print("📝 Manual Space Secrets Configuration:")
         
-        # Get tokens from environment or use provided token
-        hf_write_token = os.getenv('HF_WRITE_TOKEN', self.token)
-        hf_read_token = os.getenv('HF_READ_TOKEN', self.token)
-        hf_token = os.getenv('HF_TOKEN', self.token)  # Legacy
+        # Use the provided token as HF_TOKEN
+        hf_token = self.token
         
-        print(f"   HF_WRITE_TOKEN={hf_write_token}")
-        print(f"   HF_READ_TOKEN={hf_read_token}")
         print(f"   HF_TOKEN={hf_token}")
         
         dataset_repo = self.dataset_repo or f"{self.username}/trackio-experiments"
@@ -415,13 +391,9 @@ class TrackioSpaceDeployer:
         print(f"   TRACKIO_URL={self.space_url}")
         
         print("\n🔧 To set secrets in your Space:")
-        print("1. Go to your Space settings: {self.space_url}/settings")
+        print(f"1. Go to your Space settings: {self.space_url}/settings")
         print("2. Navigate to the 'Repository secrets' section")
         print("3. Add the following secrets:")
-        print(f"   Name: HF_WRITE_TOKEN")
-        print(f"   Value: {hf_write_token}")
-        print(f"   Name: HF_READ_TOKEN")
-        print(f"   Value: {hf_read_token}")
         print(f"   Name: HF_TOKEN")
         print(f"   Value: {hf_token}")
         print(f"   Name: TRACKIO_DATASET_REPO")
@@ -429,6 +401,7 @@ class TrackioSpaceDeployer:
         print(f"   Name: TRACKIO_URL")
         print(f"   Value: {self.space_url}")
         print("4. Save the secrets")
+        print("\nNote: HF_TOKEN starts as write token and will be switched to read token after training")
         
         return True
     

scripts/trackio_tonic/switch_to_read_token.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python3
 """
-Switch Trackio Space from Write Token to Read Token
+Switch Trackio Space HF_TOKEN from Write Token to Read Token
 
 This script switches the HF_TOKEN secret in a Trackio Space from a write token
 to a read token after the experiment is complete, for security purposes.
+The space uses only HF_TOKEN, which starts as write token and gets switched to read token.
 """
 
 import os
@@ -61,8 +62,8 @@ def switch_space_token(space_id: str, read_token: str, write_token: str) -> bool
     
     Args:
         space_id (str): The space ID (username/space-name)
-        read_token (str): The read token to set
-        write_token (str): The write token (for validation)
+        read_token (str): The read token to set as new HF_TOKEN
+        write_token (str): The write token (for authentication to update the space)
         
     Returns:
         bool: True if successful, False otherwise
@@ -93,23 +94,24 @@ def switch_space_token(space_id: str, read_token: str, write_token: str) -> bool
         # Use the write token to update the space (since we need write access)
         api = HfApi(token=write_token)
         
-        # Update the HF_TOKEN secret in the space
+        # Update the HF_TOKEN secret in the space from write token to read token
         try:
             api.add_space_secret(
                 repo_id=space_id,
                 key="HF_TOKEN",
                 value=read_token,
-                description="Hugging Face read token for dataset access (switched from write token)"
+                description="Hugging Face token for dataset access (switched from write to read for security)"
             )
-            print(f"✅ Successfully switched HF_TOKEN to read token in space: {space_id}")
+            print(f"✅ Successfully switched HF_TOKEN from write to read token in space: {space_id}")
+            print(f"🔒 Space now uses read-only permissions for enhanced security")
             return True
             
         except Exception as e:
-            print(f"❌ Failed to update space secret: {e}")
+            print(f"❌ Failed to update HF_TOKEN secret: {e}")
             return False
             
     except Exception as e:
-        print(f"❌ Error switching tokens: {e}")
+        print(f"❌ Error switching HF_TOKEN: {e}")
         return False
 
 def main():
@@ -137,12 +139,13 @@ def main():
     success = switch_space_token(space_id, read_token, write_token)
     
     if success:
-        print("\n✅ Token switch completed successfully!")
+        print("\n✅ HF_TOKEN switch completed successfully!")
         print(f"📊 Space: {space_id}")
-        print("🔒 HF_TOKEN now uses read-only permissions")
+        print("🔒 HF_TOKEN now uses read-only permissions for enhanced security")
         print("💡 The space can still read datasets but cannot write to repositories")
+        print("🎯 Training is complete - space is now secure for monitoring")
     else:
-        print("\n❌ Token switch failed!")
+        print("\n❌ HF_TOKEN switch failed!")
         print("Please check your tokens and try again.")
         sys.exit(1)
 

src/dataset_utils.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""
+Dataset utilities for Trackio experiment data management
+Provides functions for safe dataset operations with data preservation
+"""
+
+import json
+import logging
+from datetime import datetime
+from typing import Dict, Any, List, Optional, Union
+from datasets import Dataset, load_dataset
+
+logger = logging.getLogger(__name__)
+
+class TrackioDatasetManager:
+    """
+    Manager class for Trackio experiment datasets with data preservation.
+    
+    This class ensures that existing experiment data is always preserved
+    when adding new experiments or updating existing ones.
+    """
+    
+    def __init__(self, dataset_repo: str, hf_token: str):
+        """
+        Initialize the dataset manager.
+        
+        Args:
+            dataset_repo (str): HF dataset repository ID (e.g., "username/dataset-name")
+            hf_token (str): Hugging Face token for authentication
+        """
+        self.dataset_repo = dataset_repo
+        self.hf_token = hf_token
+        self._validate_repo_format()
+    
+    def _validate_repo_format(self):
+        """Validate dataset repository format"""
+        if not self.dataset_repo or '/' not in self.dataset_repo:
+            raise ValueError(f"Invalid dataset repository format: {self.dataset_repo}")
+    
+    def check_dataset_exists(self) -> bool:
+        """
+        Check if the dataset repository exists and is accessible.
+        
+        Returns:
+            bool: True if dataset exists and is accessible, False otherwise
+        """
+        try:
+            load_dataset(self.dataset_repo, token=self.hf_token)
+            logger.info(f"✅ Dataset {self.dataset_repo} exists and is accessible")
+            return True
+        except Exception as e:
+            logger.info(f"📊 Dataset {self.dataset_repo} doesn't exist or isn't accessible: {e}")
+            return False
+    
+    def load_existing_experiments(self) -> List[Dict[str, Any]]:
+        """
+        Load all existing experiments from the dataset.
+        
+        Returns:
+            List[Dict[str, Any]]: List of existing experiment dictionaries
+        """
+        try:
+            if not self.check_dataset_exists():
+                logger.info("📊 No existing dataset found, returning empty list")
+                return []
+            
+            dataset = load_dataset(self.dataset_repo, token=self.hf_token)
+            
+            if 'train' not in dataset:
+                logger.info("📊 No 'train' split found in dataset")
+                return []
+            
+            experiments = list(dataset['train'])
+            logger.info(f"📊 Loaded {len(experiments)} existing experiments")
+            
+            # Validate experiment structure
+            valid_experiments = []
+            for exp in experiments:
+                if self._validate_experiment_structure(exp):
+                    valid_experiments.append(exp)
+                else:
+                    logger.warning(f"⚠️ Skipping invalid experiment: {exp.get('experiment_id', 'unknown')}")
+            
+            logger.info(f"📊 {len(valid_experiments)} valid experiments loaded")
+            return valid_experiments
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to load existing experiments: {e}")
+            return []
+    
+    def _validate_experiment_structure(self, experiment: Dict[str, Any]) -> bool:
+        """
+        Validate that an experiment has the required structure.
+        
+        Args:
+            experiment (Dict[str, Any]): Experiment dictionary to validate
+            
+        Returns:
+            bool: True if experiment structure is valid
+        """
+        required_fields = [
+            'experiment_id', 'name', 'description', 'created_at', 
+            'status', 'metrics', 'parameters', 'artifacts', 'logs'
+        ]
+        
+        for field in required_fields:
+            if field not in experiment:
+                logger.warning(f"⚠️ Missing required field '{field}' in experiment")
+                return False
+        
+        # Validate JSON fields
+        json_fields = ['metrics', 'parameters', 'artifacts', 'logs']
+        for field in json_fields:
+            if isinstance(experiment[field], str):
+                try:
+                    json.loads(experiment[field])
+                except json.JSONDecodeError:
+                    logger.warning(f"⚠️ Invalid JSON in field '{field}' for experiment {experiment.get('experiment_id')}")
+                    return False
+        
+        return True
+    
+    def save_experiments(self, experiments: List[Dict[str, Any]], commit_message: Optional[str] = None) -> bool:
+        """
+        Save a list of experiments to the dataset, preserving data integrity.
+        
+        Args:
+            experiments (List[Dict[str, Any]]): List of experiment dictionaries
+            commit_message (Optional[str]): Custom commit message
+            
+        Returns:
+            bool: True if save was successful, False otherwise
+        """
+        try:
+            if not experiments:
+                logger.warning("⚠️ No experiments to save")
+                return False
+            
+            # Validate all experiments before saving
+            valid_experiments = []
+            for exp in experiments:
+                if self._validate_experiment_structure(exp):
+                    # Ensure last_updated is set
+                    if 'last_updated' not in exp:
+                        exp['last_updated'] = datetime.now().isoformat()
+                    valid_experiments.append(exp)
+                else:
+                    logger.error(f"❌ Invalid experiment structure: {exp.get('experiment_id', 'unknown')}")
+                    return False
+            
+            # Create dataset
+            dataset = Dataset.from_list(valid_experiments)
+            
+            # Generate commit message if not provided
+            if not commit_message:
+                commit_message = f"Update dataset with {len(valid_experiments)} experiments ({datetime.now().isoformat()})"
+            
+            # Push to hub
+            dataset.push_to_hub(
+                self.dataset_repo,
+                token=self.hf_token,
+                private=True,
+                commit_message=commit_message
+            )
+            
+            logger.info(f"✅ Successfully saved {len(valid_experiments)} experiments to {self.dataset_repo}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to save experiments to dataset: {e}")
+            return False
+    
+    def upsert_experiment(self, experiment: Dict[str, Any]) -> bool:
+        """
+        Insert a new experiment or update an existing one, preserving all other data.
+        
+        Args:
+            experiment (Dict[str, Any]): Experiment dictionary to upsert
+            
+        Returns:
+            bool: True if operation was successful, False otherwise
+        """
+        try:
+            # Validate the experiment structure
+            if not self._validate_experiment_structure(experiment):
+                logger.error(f"❌ Invalid experiment structure for {experiment.get('experiment_id', 'unknown')}")
+                return False
+            
+            # Load existing experiments
+            existing_experiments = self.load_existing_experiments()
+            
+            # Find if experiment already exists
+            experiment_id = experiment['experiment_id']
+            experiment_found = False
+            updated_experiments = []
+            
+            for existing_exp in existing_experiments:
+                if existing_exp.get('experiment_id') == experiment_id:
+                    # Update existing experiment
+                    logger.info(f"🔄 Updating existing experiment: {experiment_id}")
+                    experiment['last_updated'] = datetime.now().isoformat()
+                    updated_experiments.append(experiment)
+                    experiment_found = True
+                else:
+                    # Preserve existing experiment
+                    updated_experiments.append(existing_exp)
+            
+            # If experiment doesn't exist, add it
+            if not experiment_found:
+                logger.info(f"➕ Adding new experiment: {experiment_id}")
+                experiment['last_updated'] = datetime.now().isoformat()
+                updated_experiments.append(experiment)
+            
+            # Save all experiments
+            commit_message = f"{'Update' if experiment_found else 'Add'} experiment {experiment_id} (preserving {len(existing_experiments)} existing experiments)"
+            
+            return self.save_experiments(updated_experiments, commit_message)
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to upsert experiment: {e}")
+            return False
+    
+    def get_experiment_by_id(self, experiment_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve a specific experiment by its ID.
+        
+        Args:
+            experiment_id (str): The experiment ID to search for
+            
+        Returns:
+            Optional[Dict[str, Any]]: The experiment dictionary if found, None otherwise
+        """
+        try:
+            experiments = self.load_existing_experiments()
+            
+            for exp in experiments:
+                if exp.get('experiment_id') == experiment_id:
+                    logger.info(f"✅ Found experiment: {experiment_id}")
+                    return exp
+            
+            logger.info(f"📊 Experiment not found: {experiment_id}")
+            return None
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to get experiment {experiment_id}: {e}")
+            return None
+    
+    def list_experiments(self, status_filter: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        List all experiments, optionally filtered by status.
+        
+        Args:
+            status_filter (Optional[str]): Filter by experiment status (running, completed, failed, paused)
+            
+        Returns:
+            List[Dict[str, Any]]: List of experiments matching the filter
+        """
+        try:
+            experiments = self.load_existing_experiments()
+            
+            if status_filter:
+                filtered_experiments = [exp for exp in experiments if exp.get('status') == status_filter]
+                logger.info(f"📊 Found {len(filtered_experiments)} experiments with status '{status_filter}'")
+                return filtered_experiments
+            
+            logger.info(f"📊 Found {len(experiments)} total experiments")
+            return experiments
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to list experiments: {e}")
+            return []
+    
+    def backup_dataset(self, backup_suffix: Optional[str] = None) -> str:
+        """
+        Create a backup of the current dataset.
+        
+        Args:
+            backup_suffix (Optional[str]): Optional suffix for backup repo name
+            
+        Returns:
+            str: Backup repository name if successful, empty string otherwise
+        """
+        try:
+            if not backup_suffix:
+                backup_suffix = datetime.now().strftime('%Y%m%d_%H%M%S')
+            
+            backup_repo = f"{self.dataset_repo}-backup-{backup_suffix}"
+            
+            # Load current experiments
+            experiments = self.load_existing_experiments()
+            
+            if not experiments:
+                logger.warning("⚠️ No experiments to backup")
+                return ""
+            
+            # Create backup dataset manager
+            backup_manager = TrackioDatasetManager(backup_repo, self.hf_token)
+            
+            # Save to backup
+            success = backup_manager.save_experiments(
+                experiments, 
+                f"Backup of {self.dataset_repo} created on {datetime.now().isoformat()}"
+            )
+            
+            if success:
+                logger.info(f"✅ Backup created: {backup_repo}")
+                return backup_repo
+            else:
+                logger.error("❌ Failed to create backup")
+                return ""
+                
+        except Exception as e:
+            logger.error(f"❌ Failed to create backup: {e}")
+            return ""
+
+
+def create_dataset_manager(dataset_repo: str, hf_token: str) -> TrackioDatasetManager:
+    """
+    Factory function to create a TrackioDatasetManager instance.
+    
+    Args:
+        dataset_repo (str): HF dataset repository ID
+        hf_token (str): Hugging Face token
+        
+    Returns:
+        TrackioDatasetManager: Configured dataset manager instance
+    """
+    return TrackioDatasetManager(dataset_repo, hf_token)

src/monitoring.py

@@ -16,6 +16,7 @@ try:
     from scripts.trackio_tonic.trackio_api_client import TrackioAPIClient
     TRACKIO_AVAILABLE = True
 except ImportError:
+    TrackioAPIClient = None
     TRACKIO_AVAILABLE = False
     print("Warning: Trackio API client not available. Install with: pip install requests")
 
@@ -87,20 +88,33 @@ class SmolLM3Monitor:
         try:
             from datasets import Dataset
             from huggingface_hub import HfApi
+            try:
+                from .dataset_utils import create_dataset_manager
+            except ImportError:
+                # Try importing from same directory
+                import sys
+                import os
+                sys.path.insert(0, os.path.dirname(__file__))
+                from dataset_utils import create_dataset_manager
             
             self.hf_dataset_client = {
                 'Dataset': Dataset,
                 'HfApi': HfApi,
                 'api': HfApi(token=self.hf_token)
             }
-            logger.info("✅ HF Datasets client initialized for %s", self.dataset_repo)
+            
+            # Initialize dataset manager for safe operations
+            self.dataset_manager = create_dataset_manager(self.dataset_repo, self.hf_token)
+            logger.info("✅ HF Datasets client and manager initialized for %s", self.dataset_repo)
             
         except ImportError:
             logger.warning("⚠️ datasets or huggingface-hub not available. Install with: pip install datasets huggingface-hub")
             self.hf_dataset_client = None
+            self.dataset_manager = None
         except Exception as e:
             logger.error("Failed to initialize HF Datasets client: %s", e)
             self.hf_dataset_client = None
+            self.dataset_manager = None
     
     def _setup_trackio(self, trackio_url: Optional[str], trackio_token: Optional[str]):
         """Setup Trackio API client"""
@@ -184,55 +198,38 @@ class SmolLM3Monitor:
             self.experiment_id = f"exp_{timestamp}"
     
     def _save_to_hf_dataset(self, experiment_data: Dict[str, Any]):
-        """Save experiment data to HF Dataset"""
-        if not self.hf_dataset_client or not self.dataset_repo:
-            logger.warning("⚠️ HF Datasets not available or dataset repo not set")
+        """Save experiment data to HF Dataset with data preservation using dataset manager"""
+        if not self.dataset_manager:
+            logger.warning("⚠️ Dataset manager not available")
             return False
         
         try:
-            # Ensure dataset repository is not empty
-            if not self.dataset_repo or self.dataset_repo.strip() == '':
-                logger.error("❌ Dataset repository is empty")
-                return False
-            
-            # Validate dataset repository format
-            if '/' not in self.dataset_repo:
-                logger.error(f"❌ Invalid dataset repository format: {self.dataset_repo}")
-                return False
-            
-            Dataset = self.hf_dataset_client['Dataset']
-            api = self.hf_dataset_client['api']
-            
-            # Create dataset from experiment data with correct structure
-            # Match the structure used in setup_hf_dataset.py
-            dataset_data = [{
+            # Prepare current experiment data with standardized structure
+            current_experiment = {
                 'experiment_id': self.experiment_id or f"exp_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
                 'name': self.experiment_name,
                 'description': "SmolLM3 fine-tuning experiment",
                 'created_at': self.start_time.isoformat(),
                 'status': 'running',
-                'metrics': json.dumps(self.metrics_history),
-                'parameters': json.dumps(experiment_data),
-                'artifacts': json.dumps(self.artifacts),
-                'logs': json.dumps([]),
+                'metrics': json.dumps(self.metrics_history, default=str),
+                'parameters': json.dumps(experiment_data, default=str),
+                'artifacts': json.dumps(self.artifacts, default=str),
+                'logs': json.dumps([], default=str),
                 'last_updated': datetime.now().isoformat()
-            }]
-            
-            # Create dataset from the experiment data
-            dataset = Dataset.from_list(dataset_data)
-            
-            # Push to hub
-            dataset.push_to_hub(
-                self.dataset_repo,
-                token=self.hf_token,
-                private=True
-            )
+            }
             
-            logger.info(f"✅ Experiment data saved to HF Dataset: {self.dataset_repo}")
-            return True
+            # Use dataset manager to safely upsert the experiment
+            success = self.dataset_manager.upsert_experiment(current_experiment)
             
+            if success:
+                logger.info(f"✅ Experiment data saved to HF Dataset: {self.dataset_repo}")
+                return True
+            else:
+                logger.error(f"❌ Failed to save experiment data to HF Dataset")
+                return False
+                
         except Exception as e:
-            logger.error(f"Failed to save to HF Dataset: {e}")
+            logger.error(f"❌ Failed to save to HF Dataset: {e}")
             return False
     
     def log_configuration(self, config: Dict[str, Any]):
@@ -556,25 +553,50 @@ class SmolLM3Monitor:
             return "{}?tab=view_experiments".format(self.trackio_client.space_url)
         return None
     
-    def close(self):
-        """Close the monitoring session"""
+    def close(self, final_status: str = "completed"):
+        """
+        Close the monitoring session with final status update
+        
+        Args:
+            final_status (str): Final status for the experiment (completed, failed, etc.)
+        """
+        logger.info(f"🔚 Closing monitoring session with status: {final_status}")
+        
         if self.enable_tracking and self.trackio_client:
             try:
-                # Mark experiment as completed
+                # Mark experiment as completed in Trackio
                 result = self.trackio_client.update_experiment_status(
                     experiment_id=self.experiment_id,
-                    status="completed"
+                    status=final_status
                 )
                 if "success" in result:
-                    logger.info("Monitoring session closed")
+                    logger.info("✅ Trackio monitoring session closed")
                 else:
-                    logger.error("Failed to close monitoring session: %s", result)
+                    logger.error("❌ Failed to close Trackio monitoring session: %s", result)
+            except Exception as e:
+                logger.error("❌ Failed to close Trackio monitoring session: %s", e)
+        
+        # Final save to HF Dataset with proper status update
+        if self.dataset_manager:
+            try:
+                # Update experiment with final status
+                final_experiment_data = {
+                    'status': final_status,
+                    'experiment_end_time': datetime.now().isoformat(),
+                    'final_metrics_count': len(self.metrics_history),
+                    'total_artifacts': len(self.artifacts)
+                }
+                
+                success = self._save_to_hf_dataset(final_experiment_data)
+                if success:
+                    logger.info("✅ Final experiment data saved to HF Dataset")
+                else:
+                    logger.error("❌ Failed to save final experiment data")
+                    
             except Exception as e:
-                logger.error("Failed to close monitoring session: %s", e)
+                logger.error(f"❌ Failed to save final experiment data: {e}")
         
-        # Final save to HF Dataset
-        if self.hf_dataset_client:
-            self._save_to_hf_dataset({'status': 'completed'})
+        logger.info(f"🎯 Monitoring session closed for experiment: {self.experiment_id}")
 
 # Utility function to create monitor from config
 def create_monitor_from_config(config, experiment_name: Optional[str] = None) -> SmolLM3Monitor:

templates/spaces/demo_gpt/README.md

@@ -0,0 +1,15 @@
+---
+title: GPT-OSS-20B Multilingual Reasoner Demo
+emoji: 🌟
+colorFrom: blue
+colorTo: pink
+sdk: gradio
+sdk_version: 5.40.0
+app_file: app.py
+pinned: true
+short_description: GPT-OSS-20B Multilingual Reasoner LoRA adapter
+---
+
+This demo showcases the GPT-OSS-20B model fine-tuned with LoRA for enhanced multilingual reasoning capabilities. The model is based on OpenAI's GPT-OSS-20B base model with a LoRA adapter from Tonic.
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

templates/spaces/demo_gpt/app.py

@@ -0,0 +1,262 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline
+import torch
+from threading import Thread
+import gradio as gr
+import spaces
+import re
+import logging
+import os
+from peft import PeftModel
+
+# ----------------------------------------------------------------------
+# Environment Variables Configuration
+# ----------------------------------------------------------------------
+
+# Get model configuration from environment variables
+BASE_MODEL_ID = os.getenv('BASE_MODEL_ID', 'openai/gpt-oss-20b')
+LORA_MODEL_ID = os.getenv('LORA_MODEL_ID', os.getenv('HF_MODEL_ID', 'Tonic/gpt-oss-20b-multilingual-reasoner'))
+MODEL_NAME = os.getenv('MODEL_NAME', 'GPT-OSS Multilingual Reasoner')
+MODEL_SUBFOLDER = os.getenv('MODEL_SUBFOLDER', '')
+
+# If the LORA_MODEL_ID is the same as BASE_MODEL_ID, this is a merged model, not LoRA
+USE_LORA = LORA_MODEL_ID != BASE_MODEL_ID and not LORA_MODEL_ID.startswith(BASE_MODEL_ID)
+
+print(f"🔧 Configuration:")
+print(f"   Base Model: {BASE_MODEL_ID}")
+print(f"   Model ID: {LORA_MODEL_ID}")
+print(f"   Model Name: {MODEL_NAME}")
+print(f"   Model Subfolder: {MODEL_SUBFOLDER}")
+print(f"   Use LoRA: {USE_LORA}")
+
+# ----------------------------------------------------------------------
+# KaTeX delimiter config for Gradio
+# ----------------------------------------------------------------------
+
+LATEX_DELIMS = [
+    {"left": "$$",  "right": "$$",  "display": True},
+    {"left": "$",   "right": "$",   "display": False},
+    {"left": "\\[", "right": "\\]", "display": True},
+    {"left": "\\(", "right": "\\)", "display": False},
+]
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+
+# Load the model
+try:
+    if USE_LORA:
+        # Load base model and LoRA adapter separately
+        print(f"🔄 Loading base model: {BASE_MODEL_ID}")
+        base_model = AutoModelForCausalLM.from_pretrained(
+            BASE_MODEL_ID,
+            torch_dtype="auto",
+            device_map="auto",
+            attn_implementation="kernels-community/vllm-flash-attn3"
+        )
+        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+        
+        # Load the LoRA adapter
+        try:
+            print(f"🔄 Loading LoRA adapter: {LORA_MODEL_ID}")
+            if MODEL_SUBFOLDER and MODEL_SUBFOLDER.strip():
+                model = PeftModel.from_pretrained(base_model, LORA_MODEL_ID, subfolder=MODEL_SUBFOLDER)
+            else:
+                model = PeftModel.from_pretrained(base_model, LORA_MODEL_ID)
+            print("✅ LoRA model loaded successfully!")
+        except Exception as lora_error:
+            print(f"⚠️ LoRA adapter failed to load: {lora_error}")
+            print("🔄 Falling back to base model...")
+            model = base_model
+    else:
+        # Load merged/fine-tuned model directly
+        print(f"🔄 Loading merged model: {LORA_MODEL_ID}")
+        model_kwargs = {
+            "torch_dtype": "auto",
+            "device_map": "auto",
+            "attn_implementation": "kernels-community/vllm-flash-attn3"
+        }
+        
+        if MODEL_SUBFOLDER and MODEL_SUBFOLDER.strip():
+            model = AutoModelForCausalLM.from_pretrained(LORA_MODEL_ID, subfolder=MODEL_SUBFOLDER, **model_kwargs)
+            tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL_ID, subfolder=MODEL_SUBFOLDER)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(LORA_MODEL_ID, **model_kwargs)
+            tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL_ID)
+        print("✅ Merged model loaded successfully!")
+        
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    raise e
+
+def format_conversation_history(chat_history):
+    messages = []
+    for item in chat_history:
+        role = item["role"]
+        content = item["content"]
+        if isinstance(content, list):
+            content = content[0]["text"] if content and "text" in content[0] else str(content)
+        messages.append({"role": role, "content": content})
+    return messages
+
+def format_analysis_response(text):
+    """Enhanced response formatting with better structure and LaTeX support."""
+    # Look for analysis section followed by final response
+    m = re.search(r"analysis(.*?)assistantfinal", text, re.DOTALL | re.IGNORECASE)
+    if m:
+        reasoning = m.group(1).strip()
+        response = text.split("assistantfinal", 1)[-1].strip()
+        
+        # Clean up the reasoning section
+        reasoning = re.sub(r'^analysis\s*', '', reasoning, flags=re.IGNORECASE).strip()
+        
+        # Format with improved structure
+        formatted = (
+            f"**🤔 Analysis & Reasoning:**\n\n"
+            f"*{reasoning}*\n\n"
+            f"---\n\n"
+            f"**💬 Final Response:**\n\n{response}"
+        )
+        
+        # Ensure LaTeX delimiters are balanced
+        if formatted.count("$") % 2:
+            formatted += "$"
+            
+        return formatted
+    
+    # Fallback: clean up the text and return as-is
+    cleaned = re.sub(r'^analysis\s*', '', text, flags=re.IGNORECASE).strip()
+    if cleaned.count("$") % 2:
+        cleaned += "$"
+    return cleaned
+
+@spaces.GPU(duration=60)
+def generate_response(input_data, chat_history, max_new_tokens, system_prompt, temperature, top_p, top_k, repetition_penalty):
+    if not input_data.strip():
+        yield "Please enter a prompt."
+        return
+        
+    # Log the request
+    logging.info(f"[User] {input_data}")
+    logging.info(f"[System] {system_prompt} | Temp={temperature} | Max tokens={max_new_tokens}")
+    
+    new_message = {"role": "user", "content": input_data}
+    system_message = [{"role": "system", "content": system_prompt}] if system_prompt else []
+    processed_history = format_conversation_history(chat_history)
+    messages = system_message + processed_history + [new_message]
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    # Create streamer for proper streaming
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    # Prepare generation kwargs
+    generation_kwargs = {
+        "max_new_tokens": max_new_tokens,
+        "do_sample": True,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "repetition_penalty": repetition_penalty,
+        "pad_token_id": tokenizer.eos_token_id,
+        "streamer": streamer,
+        "use_cache": True
+    }
+    
+    # Tokenize input using the chat template
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Start generation in a separate thread
+    thread = Thread(target=model.generate, kwargs={**inputs, **generation_kwargs})
+    thread.start()
+    
+    # Stream the response with enhanced formatting
+    collected_text = ""
+    buffer = ""
+    yielded_once = False
+    
+    try:
+        for chunk in streamer:
+            if not chunk:
+                continue
+                
+            collected_text += chunk
+            buffer += chunk
+            
+            # Initial yield to show immediate response
+            if not yielded_once:
+                yield chunk
+                buffer = ""
+                yielded_once = True
+                continue
+            
+            # Yield accumulated text periodically for smooth streaming
+            if "\n" in buffer or len(buffer) > 150:
+                # Use enhanced formatting for partial text
+                partial_formatted = format_analysis_response(collected_text)
+                yield partial_formatted
+                buffer = ""
+        
+        # Final formatting with complete text
+        final_formatted = format_analysis_response(collected_text)
+        yield final_formatted
+        
+    except Exception as e:
+        logging.exception("Generation streaming failed")
+        yield f"❌ Error during generation: {e}"
+
+demo = gr.ChatInterface(
+    fn=generate_response,
+    additional_inputs=[
+        gr.Slider(label="Max new tokens", minimum=64, maximum=4096, step=1, value=2048),
+        gr.Textbox(
+            label="System Prompt",
+            value="You are a helpful assistant. Reasoning: medium",
+            lines=4,
+            placeholder="Change system prompt"
+        ),
+        gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.7),
+        gr.Slider(label="Top-p", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
+        gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=50),
+        gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.0)
+    ],
+    examples=[
+        [{"text": "Explain Newton's laws clearly and concisely with mathematical formulas"}],
+        [{"text": "Write a Python function to calculate the Fibonacci sequence"}],
+        [{"text": "What are the benefits of open weight AI models? Include analysis."}],
+        [{"text": "Solve this equation: $x^2 + 5x + 6 = 0$"}],
+    ],
+    cache_examples=False,
+    type="messages",
+    description=f"""
+
+# 🙋🏻‍♂️Welcome to 🌟{MODEL_NAME} Demo !
+
+**Model**: `{LORA_MODEL_ID}`  
+**Base**: `{BASE_MODEL_ID}`
+
+✨ **Enhanced Features:**
+- 🧠 **Advanced Reasoning**: Detailed analysis and step-by-step thinking
+- 📊 **LaTeX Support**: Mathematical formulas rendered beautifully (use `$` or `$$`)
+- 🎯 **Improved Formatting**: Clear separation of reasoning and final responses
+- 📝 **Smart Logging**: Better error handling and request tracking
+
+💡 **Usage Tips:**
+- Adjust reasoning level in system prompt (e.g., "Reasoning: high")
+- Use LaTeX for math: `$E = mc^2$` or `$$\\int x^2 dx$$`
+- Wait a couple of seconds initially for model loading
+    """,
+    fill_height=True,
+    textbox=gr.Textbox(
+        label="Query Input",
+        placeholder="Type your prompt (supports LaTeX: $x^2 + y^2 = z^2$)"
+    ),
+    stop_btn="Stop Generation",
+    multimodal=False,
+    theme=gr.themes.Soft()
+)
+
+if __name__ == "__main__":
+    demo.launch(share=True)

templates/spaces/demo_gpt/requirements.txt

@@ -0,0 +1,9 @@
+torch 
+git+https://github.com/huggingface/transformers.git
+peft
+trl
+bitsandbytes
+triton 
+accelerate
+kernels
+openai-harmony

templates/spaces/{app.py → trackio/app.py}

@@ -14,6 +14,8 @@ import plotly.graph_objects as go
 import plotly.express as px
 import pandas as pd
 import numpy as np
+import plotly.io as pio
+pio.templates.default = "plotly_white"
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -27,9 +29,24 @@ class TrackioSpace:
         self.current_experiment = None
         
         # Get dataset repository and HF token from parameters or environment variables
-        self.dataset_repo = dataset_repo or os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
+        self.dataset_repo = dataset_repo or os.environ.get('TRACKIO_DATASET_REPO', 'Tonic/trackio-experiments')
         self.hf_token = hf_token or os.environ.get('HF_TOKEN')
         
+        # Initialize dataset manager for safe operations
+        self.dataset_manager = None
+        if self.hf_token and self.dataset_repo:
+            try:
+                # Import dataset manager
+                import sys
+                sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
+                from dataset_utils import TrackioDatasetManager
+                self.dataset_manager = TrackioDatasetManager(self.dataset_repo, self.hf_token)
+                logger.info("✅ Dataset manager initialized for safe operations")
+            except ImportError:
+                logger.warning("⚠️ Dataset manager not available, using legacy data handling")
+            except Exception as e:
+                logger.warning(f"⚠️ Failed to initialize dataset manager: {e}")
+        
         logger.info(f"🔧 Using dataset repository: {self.dataset_repo}")
         
         if not self.hf_token:
@@ -38,47 +55,139 @@ class TrackioSpace:
         self._load_experiments()
         
     def _load_experiments(self):
-        """Load experiments from HF Dataset"""
+        """Load experiments from HF Dataset with data preservation support"""
         try:
+            # Try using dataset manager first for safe operations
+            if self.dataset_manager:
+                logger.info("🔍 Loading experiments using dataset manager")
+                experiments_list = self.dataset_manager.load_existing_experiments()
+                
+                # Convert list to dict format expected by the interface
+                self.experiments = {}
+                for exp_data in experiments_list:
+                    exp_id = exp_data.get('experiment_id')
+                    if exp_id:
+                        converted_experiment = self._convert_dataset_row_to_experiment(exp_data)
+                        if converted_experiment:
+                            self.experiments[exp_id] = converted_experiment
+                
+                logger.info(f"✅ Loaded {len(self.experiments)} experiments using dataset manager")
+                
+                # Sort experiments by creation date (newest first)
+                self.experiments = dict(sorted(
+                    self.experiments.items(), 
+                    key=lambda x: x[1].get('created_at', ''), 
+                    reverse=True
+                ))
+                
+                # If no experiments found, use backup
+                if not self.experiments:
+                    logger.info("📊 No experiments found in dataset, using backup data")
+                    self._load_backup_experiments()
+                
+                return
+            
+            # Fallback to direct dataset loading if dataset manager not available
             if self.hf_token:
-                from datasets import load_dataset
+                success = self._load_experiments_direct()
+                if success:
+                    return
+            
+            # Final fallback to backup data
+            logger.info("🔄 Using backup data")
+            self._load_backup_experiments()
                 
-                # Try to load the dataset
-                try:
-                    dataset = load_dataset(self.dataset_repo, token=self.hf_token)
-                    logger.info(f"✅ Loaded experiments from {self.dataset_repo}")
+        except Exception as e:
+            logger.error(f"❌ Failed to load experiments: {e}")
+            self._load_backup_experiments()
+    
+    def _load_experiments_direct(self) -> bool:
+        """Load experiments directly from HF Dataset without dataset manager"""
+        try:
+            from datasets import load_dataset
+            
+            logger.info(f"🔍 Loading experiments directly from {self.dataset_repo}")
+            dataset = load_dataset(self.dataset_repo, token=self.hf_token)
+            logger.info(f"✅ Successfully loaded dataset from {self.dataset_repo}")
+            
+            # Convert dataset to experiments dict
+            self.experiments = {}
+            if 'train' in dataset:
+                for row in dataset['train']:
+                    exp_id = row.get('experiment_id')
+                    if exp_id:
+                        converted_experiment = self._convert_dataset_row_to_experiment(row)
+                        if converted_experiment:
+                            self.experiments[exp_id] = converted_experiment
+            
+            logger.info(f"📊 Successfully loaded {len(self.experiments)} experiments from dataset")
+            
+            # Sort experiments by creation date (newest first)
+            self.experiments = dict(sorted(
+                self.experiments.items(), 
+                key=lambda x: x[1].get('created_at', ''), 
+                reverse=True
+            ))
+            
+            return True
+            
+        except Exception as e:
+            logger.warning(f"⚠️ Failed to load from dataset directly: {e}")
+            return False
+    
+    def _convert_dataset_row_to_experiment(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Convert a dataset row to experiment format, handling JSON parsing safely"""
+        try:
+            exp_id = row.get('experiment_id')
+            if not exp_id:
+                return None
+            
+            # Parse JSON fields safely
+            try:
+                metrics_raw = row.get('metrics', '[]')
+                if isinstance(metrics_raw, str):
+                    metrics = json.loads(metrics_raw) if metrics_raw else []
+                else:
+                    metrics = metrics_raw if metrics_raw else []
                     
-                    # Convert dataset to experiments dict
-                    self.experiments = {}
-                    if 'train' in dataset:
-                        for row in dataset['train']:
-                            exp_id = row.get('experiment_id')
-                            if exp_id:
-                                self.experiments[exp_id] = {
-                                    'id': exp_id,
-                                    'name': row.get('name', ''),
-                                    'description': row.get('description', ''),
-                                    'created_at': row.get('created_at', ''),
-                                    'status': row.get('status', 'running'),
-                                    'metrics': json.loads(row.get('metrics', '[]')),
-                                    'parameters': json.loads(row.get('parameters', '{}')),
-                                    'artifacts': json.loads(row.get('artifacts', '[]')),
-                                    'logs': json.loads(row.get('logs', '[]'))
-                                }
+                parameters_raw = row.get('parameters', '{}')
+                if isinstance(parameters_raw, str):
+                    parameters = json.loads(parameters_raw) if parameters_raw else {}
+                else:
+                    parameters = parameters_raw if parameters_raw else {}
                     
-                    logger.info(f"📊 Loaded {len(self.experiments)} experiments from dataset")
+                artifacts_raw = row.get('artifacts', '[]')
+                if isinstance(artifacts_raw, str):
+                    artifacts = json.loads(artifacts_raw) if artifacts_raw else []
+                else:
+                    artifacts = artifacts_raw if artifacts_raw else []
                     
-                except Exception as e:
-                    logger.warning(f"Failed to load from dataset: {e}")
-                    # Fall back to backup data
-                    self._load_backup_experiments()
-            else:
-                # No HF token, use backup data
-                self._load_backup_experiments()
-                
+                logs_raw = row.get('logs', '[]')
+                if isinstance(logs_raw, str):
+                    logs = json.loads(logs_raw) if logs_raw else []
+                else:
+                    logs = logs_raw if logs_raw else []
+                    
+            except json.JSONDecodeError as json_err:
+                logger.warning(f"JSON decode error for experiment {exp_id}: {json_err}")
+                metrics, parameters, artifacts, logs = [], {}, [], []
+            
+            return {
+                'id': exp_id,
+                'name': row.get('name', ''),
+                'description': row.get('description', ''),
+                'created_at': row.get('created_at', ''),
+                'status': row.get('status', 'running'),
+                'metrics': metrics,
+                'parameters': parameters,
+                'artifacts': artifacts,
+                'logs': logs,
+                'last_updated': row.get('last_updated', '')
+            }
+            
         except Exception as e:
-            logger.error(f"Failed to load experiments: {e}")
-            self._load_backup_experiments()
+            logger.warning(f"Failed to convert dataset row to experiment: {e}")
+            return None
     
     def _load_backup_experiments(self):
         """Load backup experiments when dataset is not available"""
@@ -312,12 +421,61 @@ class TrackioSpace:
         logger.info(f"✅ Loaded {len(backup_experiments)} backup experiments")
     
     def _save_experiments(self):
-        """Save experiments to HF Dataset"""
+        """Save experiments to HF Dataset with data preservation"""
+        try:
+            # Use dataset manager for safe operations if available
+            if self.dataset_manager:
+                logger.info("💾 Saving experiments using dataset manager (data preservation)")
+                
+                # Convert current experiments to dataset format
+                experiments_to_save = []
+                for exp_id, exp_data in self.experiments.items():
+                    experiment_entry = {
+                        'experiment_id': exp_id,
+                        'name': exp_data.get('name', ''),
+                        'description': exp_data.get('description', ''),
+                        'created_at': exp_data.get('created_at', ''),
+                        'status': exp_data.get('status', 'running'),
+                        'metrics': json.dumps(exp_data.get('metrics', []), default=str),
+                        'parameters': json.dumps(exp_data.get('parameters', {}), default=str),
+                        'artifacts': json.dumps(exp_data.get('artifacts', []), default=str),
+                        'logs': json.dumps(exp_data.get('logs', []), default=str),
+                        'last_updated': datetime.now().isoformat()
+                    }
+                    experiments_to_save.append(experiment_entry)
+                
+                # Use dataset manager to save with data preservation
+                success = self.dataset_manager.save_experiments(
+                    experiments_to_save,
+                    f"Update experiments from Trackio Space ({len(experiments_to_save)} total experiments)"
+                )
+                
+                if success:
+                    logger.info(f"✅ Successfully saved {len(experiments_to_save)} experiments with data preservation")
+                else:
+                    logger.error("❌ Failed to save experiments using dataset manager")
+                    # Fallback to legacy method
+                    self._save_experiments_legacy()
+                
+                return
+            
+            # Fallback to legacy method if dataset manager not available
+            self._save_experiments_legacy()
+                
+        except Exception as e:
+            logger.error(f"❌ Failed to save experiments: {e}")
+            # Fallback to legacy method
+            self._save_experiments_legacy()
+    
+    def _save_experiments_legacy(self):
+        """Legacy save method without data preservation (fallback only)"""
         try:
             if self.hf_token:
                 from datasets import Dataset
                 from huggingface_hub import HfApi
                 
+                logger.warning("⚠️ Using legacy save method - data preservation not guaranteed")
+                
                 # Convert experiments to dataset format
                 dataset_data = []
                 for exp_id, exp_data in self.experiments.items():
@@ -327,10 +485,10 @@ class TrackioSpace:
                         'description': exp_data.get('description', ''),
                         'created_at': exp_data.get('created_at', ''),
                         'status': exp_data.get('status', 'running'),
-                        'metrics': json.dumps(exp_data.get('metrics', [])),
-                        'parameters': json.dumps(exp_data.get('parameters', {})),
-                        'artifacts': json.dumps(exp_data.get('artifacts', [])),
-                        'logs': json.dumps(exp_data.get('logs', [])),
+                        'metrics': json.dumps(exp_data.get('metrics', []), default=str),
+                        'parameters': json.dumps(exp_data.get('parameters', {}), default=str),
+                        'artifacts': json.dumps(exp_data.get('artifacts', []), default=str),
+                        'logs': json.dumps(exp_data.get('logs', []), default=str),
                         'last_updated': datetime.now().isoformat()
                     })
                 
@@ -342,16 +500,17 @@ class TrackioSpace:
                 dataset.push_to_hub(
                     self.dataset_repo,
                     token=self.hf_token,
-                    private=True  # Make it private for security
+                    private=True,
+                    commit_message=f"Legacy update: {len(dataset_data)} experiments"
                 )
                 
-                logger.info(f"✅ Saved {len(dataset_data)} experiments to {self.dataset_repo}")
+                logger.info(f"✅ Saved {len(dataset_data)} experiments to {self.dataset_repo} (legacy method)")
                 
             else:
                 logger.warning("⚠️ No HF_TOKEN available, experiments not saved to dataset")
                 
         except Exception as e:
-            logger.error(f"Failed to save experiments to dataset: {e}")
+            logger.error(f"❌ Failed to save experiments with legacy method: {e}")
             # Fall back to local file for backup
             try:
                 data = {
@@ -363,7 +522,7 @@ class TrackioSpace:
                     json.dump(data, f, indent=2, default=str)
                 logger.info("✅ Saved backup to local file")
             except Exception as backup_e:
-                logger.error(f"Failed to save backup: {backup_e}")
+                logger.error(f"❌ Failed to save backup: {backup_e}")
     
     def create_experiment(self, name: str, description: str = "") -> Dict[str, Any]:
         """Create a new experiment"""
@@ -483,7 +642,10 @@ def update_trackio_config(hf_token: str, dataset_repo: str) -> str:
         # Reload experiments with new configuration
         trackio_space._load_experiments()
         
-        return f"✅ Configuration updated successfully!\n📊 Dataset: {trackio_space.dataset_repo}\n🔑 HF Token: {'Set' if trackio_space.hf_token else 'Not set'}\n📈 Loaded {len(trackio_space.experiments)} experiments"
+        # Check if dataset manager is available
+        manager_status = "✅ Available (data preservation enabled)" if trackio_space.dataset_manager else "⚠️ Not available (legacy mode)"
+        
+        return f"✅ Configuration updated successfully!\n📊 Dataset: {trackio_space.dataset_repo}\n🔑 HF Token: {'Set' if trackio_space.hf_token else 'Not set'}\n🛡️ Data Manager: {manager_status}\n📈 Loaded {len(trackio_space.experiments)} experiments"
         
     except Exception as e:
         return f"❌ Failed to update configuration: {str(e)}"
@@ -502,10 +664,42 @@ def test_dataset_connection(hf_token: str, dataset_repo: str) -> str:
         # Test loading the dataset
         dataset = load_dataset(dataset_repo, token=hf_token)
         
-        # Count experiments
+        # Count experiments and analyze structure
         experiment_count = len(dataset['train']) if 'train' in dataset else 0
         
-        return f"✅ Connection successful!\n📊 Dataset: {dataset_repo}\n📈 Found {experiment_count} experiments\n🔗 Dataset URL: https://huggingface.co/datasets/{dataset_repo}"
+        # Get column information
+        columns = list(dataset['train'].column_names) if 'train' in dataset else []
+        
+        # Sample first few experiment IDs
+        sample_experiments = []
+        if 'train' in dataset and experiment_count > 0:
+            for i, row in enumerate(dataset['train']):
+                if i >= 3:  # Only show first 3
+                    break
+                sample_experiments.append(row.get('experiment_id', 'unknown'))
+        
+        result = f"✅ Connection successful!\n📊 Dataset: {dataset_repo}\n📈 Found {experiment_count} experiments\n🔗 Dataset URL: https://huggingface.co/datasets/{dataset_repo}\n\n"
+        result += f"📋 Dataset Columns: {', '.join(columns)}\n"
+        if sample_experiments:
+            result += f"🔬 Sample Experiments: {', '.join(sample_experiments)}\n"
+        
+        # Test parsing one experiment if available
+        if 'train' in dataset and experiment_count > 0:
+            first_row = dataset['train'][0]
+            exp_id = first_row.get('experiment_id', 'unknown')
+            metrics_raw = first_row.get('metrics', '[]')
+            
+            try:
+                if isinstance(metrics_raw, str):
+                    metrics = json.loads(metrics_raw)
+                    metrics_count = len(metrics) if isinstance(metrics, list) else 0
+                    result += f"📊 First experiment ({exp_id}) metrics: {metrics_count} entries\n"
+                else:
+                    result += f"📊 First experiment ({exp_id}) metrics: Non-string format\n"
+            except json.JSONDecodeError as e:
+                result += f"⚠️ JSON parse error in first experiment: {e}\n"
+        
+        return result
         
     except Exception as e:
         return f"❌ Connection failed: {str(e)}\n\n💡 Troubleshooting:\n1. Check your HF token is correct\n2. Verify the dataset repository exists\n3. Ensure your token has read access to the dataset"
@@ -534,12 +728,34 @@ def create_dataset_repository(hf_token: str, dataset_repo: str) -> str:
         # Check if dataset exists
         try:
             api.dataset_info(dataset_repo)
-            return f"✅ Dataset {dataset_repo} already exists!"
+            return f"✅ Dataset {dataset_repo} already exists!\n🛡️ Data preservation is enabled for existing datasets\n🔗 View at: https://huggingface.co/datasets/{dataset_repo}"
         except:
             # Dataset doesn't exist, create it
             pass
         
-        # Create empty dataset
+        # Try to initialize dataset manager to use its repository creation
+        try:
+            # Import dataset manager
+            import sys
+            sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
+            from dataset_utils import TrackioDatasetManager
+            
+            # Create dataset manager instance
+            dataset_manager = TrackioDatasetManager(dataset_repo, hf_token)
+            
+            # Check if dataset exists using the manager
+            exists = dataset_manager.check_dataset_exists()
+            if exists:
+                return f"✅ Dataset {dataset_repo} already exists!\n🛡️ Data preservation is enabled\n🔗 View at: https://huggingface.co/datasets/{dataset_repo}"
+            
+        except ImportError:
+            # Dataset manager not available, use legacy method
+            pass
+        except Exception as e:
+            # Dataset manager failed, use legacy method
+            logger.warning(f"Dataset manager failed: {e}, using legacy method")
+        
+        # Create empty dataset with proper structure
         empty_dataset = Dataset.from_dict({
             'experiment_id': [],
             'name': [],
@@ -557,22 +773,34 @@ def create_dataset_repository(hf_token: str, dataset_repo: str) -> str:
         empty_dataset.push_to_hub(
             dataset_repo,
             token=hf_token,
-            private=True
+            private=True,
+            commit_message="Create Trackio experiment dataset with data preservation support"
         )
         
-        return f"✅ Dataset {dataset_repo} created successfully!\n🔗 View at: https://huggingface.co/datasets/{dataset_repo}\n📊 Ready to store experiments"
+        return f"✅ Dataset {dataset_repo} created successfully!\n🛡️ Data preservation is now enabled\n🔗 View at: https://huggingface.co/datasets/{dataset_repo}\n📊 Ready to store experiments safely"
         
     except Exception as e:
-        return f"❌ Failed to create dataset: {str(e)}\n\n💡 Troubleshooting:\n1. Check your HF token has write permissions\n2. Verify the username in the repository name\n3. Ensure the dataset name is valid"
+        return f"❌ Failed to create dataset: {str(e)}\n\n💡 Troubleshooting:\n1. Check your HF token has write permissions\n2. Verify the username in the repository name\n3. Ensure the dataset name is valid\n4. Check internet connectivity"
 
 # Initialize API client for remote data
 api_client = None
 try:
     from trackio_api_client import TrackioAPIClient
-    api_client = TrackioAPIClient("https://tonic-test-trackio-test.hf.space")
-    logger.info("✅ API client initialized for remote data access")
+    # Get Trackio URL from environment or use default
+    trackio_url = os.environ.get('TRACKIO_URL', 'https://tonic-test-trackio-test.hf.space')
+    
+    # Clean up URL to avoid double protocol issues
+    if trackio_url.startswith('https://https://'):
+        trackio_url = trackio_url.replace('https://https://', 'https://')
+    elif trackio_url.startswith('http://http://'):
+        trackio_url = trackio_url.replace('http://http://', 'http://')
+    
+    api_client = TrackioAPIClient(trackio_url)
+    logger.info(f"✅ API client initialized for remote data access: {trackio_url}")
 except ImportError:
     logger.warning("⚠️ API client not available, using local data only")
+except Exception as e:
+    logger.warning(f"⚠️ Failed to initialize API client: {e}, using local data only")
 
 # Add Hugging Face Spaces compatibility
 def is_huggingface_spaces():
@@ -616,6 +844,7 @@ def parse_remote_metrics_data(experiment_details: str) -> pd.DataFrame:
         lines = experiment_details.split('\n')
         metrics_data = []
         
+        # First try to parse the new format with structured experiment details
         for line in lines:
             if 'Step:' in line and 'Metrics:' in line:
                 # Extract step and metrics from the line
@@ -637,6 +866,11 @@ def parse_remote_metrics_data(experiment_details: str) -> pd.DataFrame:
                     logger.warning(f"Failed to parse metrics line: {line} - {e}")
                     continue
         
+        # If no metrics found in text format, try to parse from the dataset directly
+        if not metrics_data:
+            logger.info("No metrics found in text format, trying to parse from experiment structure")
+            # This will be handled by the updated get_remote_experiment_data function
+            
         if metrics_data:
             return pd.DataFrame(metrics_data)
         else:
@@ -647,22 +881,65 @@ def parse_remote_metrics_data(experiment_details: str) -> pd.DataFrame:
         return pd.DataFrame()
 
 def get_metrics_dataframe(experiment_id: str) -> pd.DataFrame:
-    """Get metrics as a pandas DataFrame for plotting - tries remote first, then local"""
-    # Try to get remote data first
-    remote_data = get_remote_experiment_data(experiment_id)
-    if remote_data:
-        logger.info(f"Using remote data for {experiment_id}")
-        # Parse the remote experiment details to extract metrics
-        df = parse_remote_metrics_data(remote_data["data"])
-        if not df.empty:
-            logger.info(f"Found {len(df)} metrics entries from remote data")
-            return df
-        else:
-            logger.warning(f"No metrics found in remote data for {experiment_id}")
-    
-    # Fall back to local data
-    logger.info(f"Using local data for {experiment_id}")
-    return trackio_space.get_metrics_dataframe(experiment_id)
+    """Get metrics as a pandas DataFrame for plotting - tries dataset first, then local backup"""
+    try:
+        # First try to get data directly from the dataset using the dataset manager
+        if trackio_space.dataset_manager:
+            logger.info(f"Getting metrics for {experiment_id} from dataset")
+            experiment_data = trackio_space.dataset_manager.get_experiment_by_id(experiment_id)
+            
+            if experiment_data:
+                # Parse metrics from the dataset
+                metrics_json = experiment_data.get('metrics', '[]')
+                if isinstance(metrics_json, str):
+                    try:
+                        metrics_list = json.loads(metrics_json)
+                        
+                        # Convert to DataFrame format
+                        df_data = []
+                        for metric_entry in metrics_list:
+                            if isinstance(metric_entry, dict):
+                                step = metric_entry.get('step', 0)
+                                timestamp = metric_entry.get('timestamp', '')
+                                metrics = metric_entry.get('metrics', {})
+                                
+                                row = {'step': step, 'timestamp': timestamp}
+                                row.update(metrics)
+                                df_data.append(row)
+                        
+                        if df_data:
+                            logger.info(f"Found {len(df_data)} metrics entries from dataset for {experiment_id}")
+                            return pd.DataFrame(df_data)
+                        else:
+                            logger.warning(f"No valid metrics found in dataset for {experiment_id}")
+                    except json.JSONDecodeError as e:
+                        logger.warning(f"Failed to parse metrics JSON for {experiment_id}: {e}")
+                else:
+                    logger.warning(f"Metrics data is not a JSON string for {experiment_id}")
+            else:
+                logger.warning(f"Experiment {experiment_id} not found in dataset")
+        
+        # Try legacy remote data approach
+        remote_data = get_remote_experiment_data(experiment_id)
+        if remote_data:
+            logger.info(f"Using remote API data for {experiment_id}")
+            # Parse the remote experiment details to extract metrics
+            df = parse_remote_metrics_data(remote_data["data"])
+            if not df.empty:
+                logger.info(f"Found {len(df)} metrics entries from remote API")
+                return df
+            else:
+                logger.warning(f"No metrics found in remote API data for {experiment_id}")
+        
+        # Fall back to local data
+        logger.info(f"Using local backup data for {experiment_id}")
+        return trackio_space.get_metrics_dataframe(experiment_id)
+        
+    except Exception as e:
+        logger.error(f"Error getting metrics dataframe for {experiment_id}: {e}")
+        # Fall back to local data
+        logger.info(f"Falling back to local data for {experiment_id}")
+        return trackio_space.get_metrics_dataframe(experiment_id)
 
 def create_experiment_interface(name: str, description: str) -> str:
     """Create a new experiment"""
@@ -919,12 +1196,622 @@ def create_demo_experiment():
     except Exception as e:
         return f"❌ Error creating demo experiment: {str(e)}"
 
+
+# Helper functions for the new interface
+def get_experiment_dropdown_choices() -> list:
+    """Get the list of experiments for the dropdown"""
+    experiments = list(trackio_space.experiments.keys())
+    if not experiments:
+        return ["No experiments available"]
+    return experiments
+
+def refresh_experiment_dropdown() -> tuple:
+    """Refresh the experiment dropdown and return current choices"""
+    choices = get_experiment_dropdown_choices()
+    current_value = choices[0] if choices and choices[0] != "No experiments available" else None
+    return gr.Dropdown(choices=choices, value=current_value)
+
+def get_available_metrics_for_experiments(experiment_ids: list) -> list:
+    """Get all available metrics across selected experiments"""
+    try:
+        all_metrics = set()
+        for exp_id in experiment_ids:
+            df = get_metrics_dataframe(exp_id)
+            if not df.empty:
+                # Get numeric columns (excluding step and timestamp)
+                numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
+                numeric_cols = [col for col in numeric_cols if col not in ['step']]
+                all_metrics.update(numeric_cols)
+        
+        return sorted(list(all_metrics))
+    except Exception as e:
+        logger.error(f"Error getting available metrics: {str(e)}")
+        return ["loss", "accuracy"]
+
+def create_test_plot() -> go.Figure:
+    """Create a simple test plot to verify plotly rendering works"""
+    try:
+        # Create simple test data
+        x = [1, 2, 3, 4, 5]
+        y = [1, 4, 2, 3, 5]
+        
+        fig = go.Figure()
+        fig.add_trace(go.Scatter(
+            x=x,
+            y=y,
+            mode='lines+markers',
+            name='Test Data',
+            line=dict(width=2, color='blue'),
+            marker=dict(size=5, color='red'),
+            connectgaps=True,
+            hovertemplate='<b>X:</b> %{x}<br><b>Y:</b> %{y}<extra></extra>'
+        ))
+        
+        fig.update_layout(
+            title="Test Plot - If you can see this, plotly is working!",
+            xaxis_title="X Axis",
+            yaxis_title="Y Axis",
+            plot_bgcolor='white',
+            paper_bgcolor='white',
+            font=dict(size=14),
+            margin=dict(l=50, r=50, t=80, b=50)
+        )
+        
+        fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
+        fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
+        
+        logger.info("Test plot created successfully")
+        return fig
+        
+    except Exception as e:
+        logger.error(f"Error creating test plot: {str(e)}")
+        fig = go.Figure()
+        fig.add_annotation(
+            text=f"Test plot error: {str(e)}",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5, showarrow=False,
+            font=dict(size=14, color="red")
+        )
+        return fig
+
+def get_experiment_status_summary(experiment_id: str) -> str:
+    """Get a formatted summary of experiment status and metadata"""
+    try:
+        experiment = trackio_space.get_experiment(experiment_id)
+        if not experiment:
+            return f"Experiment {experiment_id} not found."
+        
+        summary = f"📋 EXPERIMENT STATUS SUMMARY\n{'='*50}\n"
+        summary += f"ID: {experiment['id']}\n"
+        summary += f"Name: {experiment['name']}\n"
+        summary += f"Description: {experiment['description']}\n"
+        summary += f"Status: {experiment['status']}\n"
+        summary += f"Created: {experiment['created_at']}\n"
+        summary += f"Metrics entries: {len(experiment['metrics'])}\n"
+        summary += f"Parameters: {len(experiment['parameters'])}\n"
+        summary += f"Artifacts: {len(experiment['artifacts'])}\n"
+        summary += f"Logs: {len(experiment['logs'])}\n"
+        
+        # Add latest metrics if available
+        if experiment['metrics']:
+            latest = experiment['metrics'][-1]
+            summary += f"\n📈 LATEST METRICS (Step {latest.get('step', 'N/A')}):\n"
+            for k, v in latest.get('metrics', {}).items():
+                summary += f"  {k}: {v}\n"
+        
+        return summary
+    except Exception as e:
+        return f"Error generating status summary: {str(e)}"
+
+def get_experiment_parameters_summary(experiment_id: str) -> str:
+    """Get a formatted summary of experiment parameters"""
+    try:
+        experiment = trackio_space.get_experiment(experiment_id)
+        if not experiment:
+            return f"Experiment {experiment_id} not found."
+        
+        params = experiment.get('parameters', {})
+        if not params:
+            return "No parameters logged for this experiment."
+        
+        summary = f"🔧 PARAMETERS FOR {experiment_id}\n{'='*50}\n"
+        
+        # Group parameters by category
+        model_params = {k: v for k, v in params.items() if 'model' in k.lower() or 'name' in k.lower()}
+        training_params = {k: v for k, v in params.items() if any(x in k.lower() for x in ['learning', 'batch', 'epoch', 'step', 'iter', 'optimizer'])}
+        data_params = {k: v for k, v in params.items() if any(x in k.lower() for x in ['data', 'dataset', 'file', 'split'])}
+        other_params = {k: v for k, v in params.items() if k not in model_params and k not in training_params and k not in data_params}
+        
+        if model_params:
+            summary += "🤖 MODEL PARAMETERS:\n"
+            for k, v in model_params.items():
+                summary += f"  {k}: {v}\n"
+            summary += "\n"
+        
+        if training_params:
+            summary += "🏃 TRAINING PARAMETERS:\n"
+            for k, v in training_params.items():
+                summary += f"  {k}: {v}\n"
+            summary += "\n"
+        
+        if data_params:
+            summary += "📁 DATA PARAMETERS:\n"
+            for k, v in data_params.items():
+                summary += f"  {k}: {v}\n"
+            summary += "\n"
+        
+        if other_params:
+            summary += "⚙️ OTHER PARAMETERS:\n"
+            for k, v in other_params.items():
+                summary += f"  {k}: {v}\n"
+        
+        return summary
+    except Exception as e:
+        return f"Error generating parameters summary: {str(e)}"
+
+def get_experiment_metrics_summary(experiment_id: str) -> str:
+    """Get a summary of all metrics for an experiment"""
+    try:
+        df = get_metrics_dataframe(experiment_id)
+        if df.empty:
+            return "No metrics data available for this experiment.\n\n💡 This could mean:\n• The experiment hasn't started logging metrics yet\n• The experiment is using a different data format\n• No training has been performed on this experiment"
+        
+        # Get numeric columns (excluding step and timestamp)
+        numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
+        numeric_cols = [col for col in numeric_cols if col not in ['step']]
+        
+        if not numeric_cols:
+            return "No numeric metrics found for this experiment.\n\n💡 This could mean:\n• Only timestamp data is available\n• Metrics are stored in a different format\n• The experiment hasn't logged any numeric metrics yet"
+        
+        summary = f"📊 METRICS SUMMARY FOR {experiment_id}\n{'='*50}\n"
+        summary += f"Total data points: {len(df)}\n"
+        summary += f"Steps range: {df['step'].min()} - {df['step'].max()}\n"
+        summary += f"Available metrics: {', '.join(numeric_cols)}\n\n"
+        
+        for col in numeric_cols:
+            if col in df.columns:
+                values = df[col].dropna()
+                if len(values) > 0:
+                    summary += f"{col}:\n"
+                    summary += f"  Min: {values.min():.6f}\n"
+                    summary += f"  Max: {values.max():.6f}\n"
+                    summary += f"  Mean: {values.mean():.6f}\n"
+                    summary += f"  Latest: {values.iloc[-1]:.6f}\n\n"
+        
+        return summary
+    except Exception as e:
+        return f"Error generating metrics summary: {str(e)}"
+
+def create_combined_metrics_plot(experiment_id: str) -> go.Figure:
+    """Create a combined plot showing all metrics for an experiment"""
+    try:
+        if not experiment_id:
+            fig = go.Figure()
+            fig.add_annotation(
+                text="No experiment selected",
+                xref="paper", yref="paper",
+                x=0.5, y=0.5, showarrow=False,
+                font=dict(size=16, color="gray")
+            )
+            fig.update_layout(
+                title="Select an Experiment",
+                plot_bgcolor='white', paper_bgcolor='white'
+            )
+            return fig
+            
+        df = get_metrics_dataframe(experiment_id)
+        if df.empty:
+            fig = go.Figure()
+            fig.add_annotation(
+                text="No metrics data available for this experiment",
+                xref="paper", yref="paper",
+                x=0.5, y=0.5, showarrow=False,
+                font=dict(size=16, color="red")
+            )
+            fig.update_layout(
+                title="No Data Available",
+                plot_bgcolor='white', paper_bgcolor='white'
+            )
+            return fig
+        
+        # Get numeric columns (excluding step and timestamp)
+        numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
+        numeric_cols = [col for col in numeric_cols if col not in ['step']]
+        
+        if not numeric_cols:
+            fig = go.Figure()
+            fig.add_annotation(
+                text="No numeric metrics found for this experiment",
+                xref="paper", yref="paper",
+                x=0.5, y=0.5, showarrow=False,
+                font=dict(size=16, color="orange")
+            )
+            fig.update_layout(
+                title="No Metrics Found",
+                plot_bgcolor='white', paper_bgcolor='white'
+            )
+            return fig
+        
+        # Create subplots for multiple metrics
+        from plotly.subplots import make_subplots
+        
+        # Determine number of rows and columns for subplots
+        n_metrics = len(numeric_cols)
+        n_cols = min(3, n_metrics)  # Max 3 columns
+        n_rows = (n_metrics + n_cols - 1) // n_cols
+        
+        fig = make_subplots(
+            rows=n_rows, cols=n_cols,
+            subplot_titles=numeric_cols,
+            vertical_spacing=0.05,
+            horizontal_spacing=0.1
+        )
+        
+        # Define colors for different metrics
+        colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown', 'pink', 'gray', 'cyan', 'magenta']
+        
+        for i, metric in enumerate(numeric_cols):
+            if metric in df.columns and not df[metric].isna().all():
+                row = (i // n_cols) + 1
+                col = (i % n_cols) + 1
+                color = colors[i % len(colors)]
+                
+                fig.add_trace(
+                    go.Scatter(
+                        x=df['step'].tolist(),
+                        y=df[metric].tolist(),
+                        mode='lines+markers',
+                        name=metric,
+                        line=dict(width=2, color=color),
+                        marker=dict(size=4, color=color),
+                        showlegend=False,
+                        connectgaps=True
+                    ),
+                    row=row, col=col
+                )
+        
+        fig.update_layout(
+            title=f"All Metrics for Experiment {experiment_id}",
+            height=350 * n_rows,
+            plot_bgcolor='white',
+            paper_bgcolor='white',
+            font=dict(size=12),
+            margin=dict(l=50, r=50, t=80, b=50)
+        )
+        
+        # Update all subplot axes
+        for i in range(1, n_rows + 1):
+            for j in range(1, n_cols + 1):
+                fig.update_xaxes(
+                    showgrid=True, gridwidth=1, gridcolor='lightgray',
+                    zeroline=True, zerolinecolor='black',
+                    row=i, col=j
+                )
+                fig.update_yaxes(
+                    showgrid=True, gridwidth=1, gridcolor='lightgray',
+                    zeroline=True, zerolinecolor='black',
+                    row=i, col=j
+                )
+        
+        return fig
+        
+    except Exception as e:
+        logger.error(f"Error creating combined metrics plot: {str(e)}")
+        fig = go.Figure()
+        fig.add_annotation(
+            text=f"Error creating combined plot: {str(e)}",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5, showarrow=False,
+            font=dict(size=14, color="red")
+        )
+        return fig
+
+def update_dashboard(experiment_id: str) -> tuple:
+    """Update all dashboard components for a selected experiment"""
+    try:
+        if not experiment_id or experiment_id == "No experiments available":
+            return (
+                "Please select an experiment from the dropdown.",
+                "No experiment selected.",
+                "No experiment selected.",
+                create_combined_metrics_plot(""),
+                "No experiment selected."
+            )
+        
+        # Get all the dashboard components
+        status_summary = get_experiment_status_summary(experiment_id)
+        parameters_summary = get_experiment_parameters_summary(experiment_id)
+        metrics_summary = get_experiment_metrics_summary(experiment_id)
+        combined_plot = create_combined_metrics_plot(experiment_id)
+        
+        # Create a combined summary
+        combined_summary = f"{status_summary}\n\n{parameters_summary}\n\n{metrics_summary}"
+        
+        return (
+            status_summary,
+            parameters_summary,
+            metrics_summary,
+            combined_plot,
+            combined_summary
+        )
+    except Exception as e:
+        error_msg = f"Error updating dashboard: {str(e)}"
+        return (error_msg, error_msg, error_msg, create_combined_metrics_plot(""), error_msg)
+
+def update_dashboard_metric_plot(experiment_id: str, metric_name: str = "loss") -> go.Figure:
+    """Update the dashboard metric plot for a selected experiment and metric"""
+    try:
+        if not experiment_id or experiment_id == "No experiments available":
+            return create_metrics_plot("", metric_name)
+        
+        return create_metrics_plot(experiment_id, metric_name)
+    except Exception as e:
+        logger.error(f"Error updating dashboard metric plot: {str(e)}")
+        return create_metrics_plot("", metric_name)
+
+def create_experiment_comparison_from_selection(selected_experiments: list, selected_metrics: list) -> go.Figure:
+    """Create experiment comparison from checkbox selections"""
+    try:
+        if not selected_experiments:
+            fig = go.Figure()
+            fig.add_annotation(
+                text="Please select at least one experiment to compare",
+                xref="paper", yref="paper",
+                x=0.5, y=0.5, showarrow=False,
+                font=dict(size=16, color="orange")
+            )
+            fig.update_layout(
+                title="No Experiments Selected",
+                plot_bgcolor='white', paper_bgcolor='white'
+            )
+            return fig
+        
+        if not selected_metrics:
+            fig = go.Figure()
+            fig.add_annotation(
+                text="Please select at least one metric to compare",
+                xref="paper", yref="paper",
+                x=0.5, y=0.5, showarrow=False,
+                font=dict(size=16, color="orange")
+            )
+            fig.update_layout(
+                title="No Metrics Selected",
+                plot_bgcolor='white', paper_bgcolor='white'
+            )
+            return fig
+        
+        # Use the existing comparison function with comma-separated IDs
+        experiment_ids_str = ",".join(selected_experiments)
+        return create_experiment_comparison(experiment_ids_str)
+        
+    except Exception as e:
+        logger.error(f"Error creating comparison from selection: {str(e)}")
+        fig = go.Figure()
+        fig.add_annotation(
+            text=f"Error creating comparison: {str(e)}",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5, showarrow=False,
+            font=dict(size=14, color="red")
+        )
+        return fig
+
+def refresh_comparison_options() -> tuple:
+    """Refresh the experiment and metric options for comparison"""
+    try:
+        # Get updated experiment choices
+        experiment_choices = get_experiment_dropdown_choices()
+        if experiment_choices == ["No experiments available"]:
+            experiment_choices = []
+        
+        # Get available metrics from all experiments
+        all_experiments = list(trackio_space.experiments.keys())
+        available_metrics = get_available_metrics_for_experiments(all_experiments)
+        
+        # Default to common metrics if available
+        default_metrics = []
+        common_metrics = ["loss", "accuracy", "learning_rate", "gpu_memory"]
+        for metric in common_metrics:
+            if metric in available_metrics:
+                default_metrics.append(metric)
+        
+        # If no common metrics, use first few available
+        if not default_metrics and available_metrics:
+            default_metrics = available_metrics[:2]
+        
+        return gr.CheckboxGroup(choices=experiment_choices, value=[]), gr.CheckboxGroup(choices=available_metrics, value=default_metrics)
+    except Exception as e:
+        logger.error(f"Error refreshing comparison options: {str(e)}")
+        return gr.CheckboxGroup(choices=[], value=[]), gr.CheckboxGroup(choices=["loss", "accuracy"], value=[])
+
 # Create Gradio interface
 with gr.Blocks(title="Trackio - Experiment Tracking", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🚀 Trackio Experiment Tracking & Monitoring")
     gr.Markdown("Monitor and track your ML experiments with real-time visualization!")
     
     with gr.Tabs():
+        # Dashboard Tab (NEW)
+        with gr.Tab("📊 Dashboard"):
+            gr.Markdown("### Comprehensive Experiment Dashboard")
+            gr.Markdown("Select an experiment to view all its data, plots, and information in one place.")
+            
+            # Row 1: Experiment Selection
+            with gr.Row():
+                with gr.Column(scale=3):
+                    # Experiment selection dropdown
+                    experiment_dropdown = gr.Dropdown(
+                        label="Select Experiment",
+                        choices=get_experiment_dropdown_choices(),
+                        value=get_experiment_dropdown_choices()[0] if get_experiment_dropdown_choices() and get_experiment_dropdown_choices()[0] != "No experiments available" else None,
+                        info="Choose an experiment to view its dashboard"
+                    )
+                
+                with gr.Column(scale=1):
+                    with gr.Row():
+                        refresh_dropdown_btn = gr.Button("🔄 Refresh List", variant="secondary", size="sm")
+                        refresh_dashboard_btn = gr.Button("🔄 Refresh Dashboard", variant="primary", size="sm")
+            
+            # Row 2: All Metrics Plots
+            with gr.Row():
+                with gr.Column(scale=3):
+                    with gr.Row():
+                        gr.Markdown("### 📈 All Metrics Plots")
+                    with gr.Row():
+                        with gr.Column(scale=3):    
+                            dashboard_plots = gr.Plot(
+                                label="Training Metrics", 
+                                container=True, 
+                                show_label=True,
+                                elem_classes=["plot-container"]
+                            )
+                
+            # Row 3: Training Metrics Visualization Accordion
+            with gr.Row():
+                with gr.Accordion("📈 Training Metrics Visualization", open=False):
+                    with gr.Row():
+                        with gr.Column():
+                            metric_dropdown = gr.Dropdown(
+                                label="Metric to Plot",
+                                choices=[
+                                    "loss", "accuracy", "learning_rate", "gpu_memory", "training_time",
+                                    "total_tokens", "truncated_tokens", "padding_tokens", "throughput", "step_time",
+                                    "batch_size", "seq_len", "token_acc", "train/gate_ortho", "train/center"
+                                ],
+                                value="loss"
+                            )
+                            plot_btn = gr.Button("Create Plot", variant="primary")
+                            test_plot_btn = gr.Button("Test Plot Rendering", variant="secondary")
+                        
+                    with gr.Row():
+                        dashboard_metric_plot = gr.Plot(
+                            label="Training Metrics", 
+                            container=True, 
+                            show_label=True,
+                            elem_classes=["plot-container"]
+                        )
+                    
+                    plot_btn.click(
+                        create_metrics_plot,
+                        inputs=[experiment_dropdown, metric_dropdown],
+                        outputs=dashboard_metric_plot
+                    )
+                    
+                    test_plot_btn.click(
+                        create_test_plot,
+                        inputs=[],
+                        outputs=dashboard_metric_plot
+                    )
+            
+            # Row 4: Accordion with Detailed Information
+            with gr.Row():
+                with gr.Accordion("📋 Experiment Details", open=False):
+                    with gr.Tabs():
+                        with gr.Tab("📋 Status"):
+                            dashboard_status = gr.Textbox(
+                                label="Experiment Status",
+                                lines=8,
+                                interactive=False
+                            )
+                        
+                        with gr.Tab("🔧 Parameters"):
+                            dashboard_parameters = gr.Textbox(
+                                label="Experiment Parameters",
+                                lines=12,
+                                interactive=False
+                            )
+                        
+                        with gr.Tab("📊 Metrics Summary"):
+                            dashboard_metrics = gr.Textbox(
+                                label="Metrics Summary",
+                                lines=12,
+                                interactive=False
+                            )
+                        
+                        with gr.Tab("📋 Complete Summary"):
+                            dashboard_summary = gr.Textbox(
+                                label="Full Experiment Summary",
+                                lines=20,
+                                interactive=False
+                            )
+            
+            # Connect the dashboard update function
+            experiment_dropdown.change(
+                update_dashboard,
+                inputs=[experiment_dropdown],
+                outputs=[dashboard_status, dashboard_parameters, dashboard_metrics, dashboard_plots, dashboard_summary]
+            )
+            
+            refresh_dashboard_btn.click(
+                update_dashboard,
+                inputs=[experiment_dropdown],
+                outputs=[dashboard_status, dashboard_parameters, dashboard_metrics, dashboard_plots, dashboard_summary]
+            )
+            
+            # Connect the metric plot update function
+            metric_dropdown.change(
+                update_dashboard_metric_plot,
+                inputs=[experiment_dropdown, metric_dropdown],
+                outputs=[dashboard_metric_plot]
+            )
+            
+            refresh_dropdown_btn.click(
+                refresh_experiment_dropdown,
+                inputs=[],
+                outputs=[experiment_dropdown]
+            )
+        
+
+        # Experiment Comparison Tab
+        with gr.Tab("📊 Experiment Comparison"):
+            gr.Markdown("### Compare Multiple Experiments")
+            gr.Markdown("Select experiments and metrics to compare from the available options below.")
+            
+            # Selection controls
+            with gr.Row():
+                with gr.Column(scale=2):
+                    gr.Markdown("### Available Experiments")
+                    experiment_checkboxes = gr.CheckboxGroup(
+                        label="Select Experiments to Compare",
+                        choices=get_experiment_dropdown_choices(),
+                        value=[],
+                        info="Choose experiments to include in the comparison"
+                    )
+                    
+                    gr.Markdown("### Available Metrics")
+                    metric_checkboxes = gr.CheckboxGroup(
+                        label="Select Metrics to Compare",
+                        choices=get_available_metrics_for_experiments(list(trackio_space.experiments.keys())),
+                        value=["loss", "accuracy"],
+                        info="Choose metrics to include in the comparison"
+                    )
+                    
+                    with gr.Row():
+                        comparison_btn = gr.Button("Compare Selected", variant="primary")
+                        refresh_options_btn = gr.Button("🔄 Refresh Options", variant="secondary")
+                
+                with gr.Column(scale=1):
+                    gr.Markdown("### Comparison Results")
+                    gr.Markdown("The comparison will show subplots for the selected metrics across the selected experiments.")
+            
+            # Comparison plots as subplots
+            comparison_plot = gr.Plot(
+                label="Experiment Comparison Dashboard", 
+                container=True, 
+                show_label=True,
+                elem_classes=["plot-container"]
+            )
+            
+            comparison_btn.click(
+                create_experiment_comparison_from_selection,
+                inputs=[experiment_checkboxes, metric_checkboxes],
+                outputs=comparison_plot
+            )
+            
+            refresh_options_btn.click(
+                refresh_comparison_options,
+                inputs=[],
+                outputs=[experiment_checkboxes, metric_checkboxes]
+            )
+
         # Configuration Tab
         with gr.Tab("⚙️ Configuration"):
             gr.Markdown("### Configure HF Datasets Connection")
@@ -941,7 +1828,7 @@ with gr.Blocks(title="Trackio - Experiment Tracking", theme=gr.themes.Soft()) as
                     dataset_repo_input = gr.Textbox(
                         label="Dataset Repository",
                         placeholder="your-username/your-dataset-name",
-                        value="tonic/trackio-experiments",
+                        value="Tonic/trackio-experiments",
                         info="HF Dataset repository for experiment storage"
                     )
                     
@@ -953,9 +1840,9 @@ with gr.Blocks(title="Trackio - Experiment Tracking", theme=gr.themes.Soft()) as
                     gr.Markdown("### Current Configuration")
                     current_config_output = gr.Textbox(
                         label="Status",
-                        lines=8,
+                        lines=10,
                         interactive=False,
-                        value=f"📊 Dataset: {trackio_space.dataset_repo}\n🔑 HF Token: {'Set' if trackio_space.hf_token else 'Not set'}\n📈 Experiments: {len(trackio_space.experiments)}"
+                        value=f"📊 Dataset: {trackio_space.dataset_repo}\n🔑 HF Token: {'Set' if trackio_space.hf_token else 'Not set'}\n🛡️ Data Preservation: {'✅ Enabled' if trackio_space.dataset_manager else '⚠️ Legacy Mode'}\n📈 Experiments: {len(trackio_space.experiments)}\n📋 Available Experiments: {', '.join(list(trackio_space.experiments.keys())[:3])}{'...' if len(trackio_space.experiments) > 3 else ''}"
                     )
                 
                 with gr.Column():
@@ -978,12 +1865,204 @@ with gr.Blocks(title="Trackio - Experiment Tracking", theme=gr.themes.Soft()) as
                     - `HF_TOKEN`: Your Hugging Face token
                     - `TRACKIO_DATASET_REPO`: Dataset repository
                     
+                    **Data Preservation:**
+                    - ✅ **Enabled**: All experiment data is preserved when adding/updating experiments
+                    - ⚠️ **Legacy Mode**: Data preservation not guaranteed (fallback mode)
+                    - Data preservation requires the dataset management utilities to be available
+                    
                     **Actions:**
                     - **Update Configuration**: Apply new settings and reload experiments
                     - **Test Connection**: Verify access to the dataset repository
                     - **Create Dataset**: Create a new dataset repository if it doesn't exist
                     """)
             
+            # Experiment Management Accordion
+            with gr.Accordion("🔧 Experiment Management", open=False):
+                with gr.Tabs():
+                    # Create Experiment Tab
+                    with gr.Tab("Create Experiment"):
+                        gr.Markdown("### Create a New Experiment")
+                        with gr.Row():
+                            with gr.Column():
+                                create_exp_name = gr.Textbox(
+                                    label="Experiment Name",
+                                    placeholder="my_smollm3_finetune",
+                                    value="smollm3_finetune"
+                                )
+                                create_exp_description = gr.Textbox(
+                                    label="Description",
+                                    placeholder="Fine-tuning SmolLM3 model on custom dataset",
+                                    value="SmolLM3 fine-tuning experiment"
+                                )
+                                create_exp_btn = gr.Button("Create Experiment", variant="primary")
+                            
+                            with gr.Column():
+                                create_exp_output = gr.Textbox(
+                                    label="Result",
+                                    lines=5,
+                                    interactive=False
+                                )
+                        
+                        create_exp_btn.click(
+                            create_experiment_interface,
+                            inputs=[create_exp_name, create_exp_description],
+                            outputs=[create_exp_output, experiment_dropdown]
+                        )
+                    
+                    # Log Metrics Tab
+                    with gr.Tab("Log Metrics"):
+                        gr.Markdown("### Log Training Metrics")
+                        with gr.Row():
+                            with gr.Column():
+                                log_metrics_exp_id = gr.Textbox(
+                                    label="Experiment ID",
+                                    placeholder="exp_20231201_143022"
+                                )
+                                log_metrics_json = gr.Textbox(
+                                    label="Metrics (JSON)",
+                                    placeholder='{"loss": 0.5, "accuracy": 0.85, "learning_rate": 2e-5}',
+                                    value='{"loss": 0.5, "accuracy": 0.85, "learning_rate": 2e-5, "gpu_memory": 22.5}'
+                                )
+                                log_metrics_step = gr.Textbox(
+                                    label="Step (optional)",
+                                    placeholder="100"
+                                )
+                                log_metrics_btn = gr.Button("Log Metrics", variant="primary")
+                            
+                            with gr.Column():
+                                log_metrics_output = gr.Textbox(
+                                    label="Result",
+                                    lines=5,
+                                    interactive=False
+                                )
+                        
+                        log_metrics_btn.click(
+                            log_metrics_interface,
+                            inputs=[log_metrics_exp_id, log_metrics_json, log_metrics_step],
+                            outputs=log_metrics_output
+                        )
+                    
+                    # Log Parameters Tab
+                    with gr.Tab("Log Parameters"):
+                        gr.Markdown("### Log Experiment Parameters")
+                        with gr.Row():
+                            with gr.Column():
+                                log_params_exp_id = gr.Textbox(
+                                    label="Experiment ID",
+                                    placeholder="exp_20231201_143022"
+                                )
+                                log_params_json = gr.Textbox(
+                                    label="Parameters (JSON)",
+                                    placeholder='{"learning_rate": 2e-5, "batch_size": 4}',
+                                    value='{"learning_rate": 3.5e-6, "batch_size": 8, "model_name": "HuggingFaceTB/SmolLM3-3B", "max_iters": 18000, "mixed_precision": "bf16"}'
+                                )
+                                log_params_btn = gr.Button("Log Parameters", variant="primary")
+                            
+                            with gr.Column():
+                                log_params_output = gr.Textbox(
+                                    label="Result",
+                                    lines=5,
+                                    interactive=False
+                                )
+                        
+                        log_params_btn.click(
+                            log_parameters_interface,
+                            inputs=[log_params_exp_id, log_params_json],
+                            outputs=log_params_output
+                        )
+                    
+                    # View Experiments Tab
+                    with gr.Tab("View Experiments"):
+                        gr.Markdown("### View Experiment Details")
+                        with gr.Row():
+                            with gr.Column():
+                                view_exp_id = gr.Textbox(
+                                    label="Experiment ID",
+                                    placeholder="exp_20231201_143022"
+                                )
+                                view_btn = gr.Button("View Experiment", variant="primary")
+                                list_btn = gr.Button("List All Experiments", variant="secondary")
+                            
+                            with gr.Column():
+                                view_output = gr.Textbox(
+                                    label="Experiment Details",
+                                    lines=20,
+                                    interactive=False
+                                )
+                        
+                        view_btn.click(
+                            get_experiment_details,
+                            inputs=[view_exp_id],
+                            outputs=view_output
+                        )
+                        
+                        list_btn.click(
+                            list_experiments_interface,
+                            inputs=[],
+                            outputs=view_output
+                        )
+                    
+                    # Update Status Tab
+                    with gr.Tab("Update Status"):
+                        gr.Markdown("### Update Experiment Status")
+                        with gr.Row():
+                            with gr.Column():
+                                status_exp_id = gr.Textbox(
+                                    label="Experiment ID",
+                                    placeholder="exp_20231201_143022"
+                                )
+                                status_dropdown = gr.Dropdown(
+                                    label="Status",
+                                    choices=["running", "completed", "failed", "paused"],
+                                    value="running"
+                                )
+                                update_status_btn = gr.Button("Update Status", variant="primary")
+                            
+                            with gr.Column():
+                                status_output = gr.Textbox(
+                                    label="Result",
+                                    lines=3,
+                                    interactive=False
+                                )
+                        
+                        update_status_btn.click(
+                            update_experiment_status_interface,
+                            inputs=[status_exp_id, status_dropdown],
+                            outputs=status_output
+                        )
+                    
+                    # Demo Data Tab
+                    with gr.Tab("Demo Data"):
+                        gr.Markdown("### Generate Demo Training Data")
+                        gr.Markdown("Use this to simulate training data for testing the interface")
+                        with gr.Row():
+                            with gr.Column():
+                                demo_exp_id = gr.Textbox(
+                                    label="Experiment ID",
+                                    placeholder="exp_20231201_143022"
+                                )
+                                demo_btn = gr.Button("Generate Demo Data", variant="primary")
+                                create_demo_btn = gr.Button("Create Demo Experiment", variant="secondary")
+                            
+                            with gr.Column():
+                                demo_output = gr.Textbox(
+                                    label="Result",
+                                    lines=5,
+                                    interactive=False
+                                )
+                        
+                        demo_btn.click(
+                            simulate_training_data,
+                            inputs=[demo_exp_id],
+                            outputs=[demo_output, dashboard_status, dashboard_parameters, dashboard_metrics, dashboard_plots, dashboard_summary]
+                        )
+                        
+                        create_demo_btn.click(
+                            create_demo_experiment,
+                            inputs=[],
+                            outputs=[demo_output, experiment_dropdown]
+                        )
+            
             update_config_btn.click(
                 update_trackio_config,
                 inputs=[hf_token_input, dataset_repo_input],
@@ -1001,237 +2080,9 @@ with gr.Blocks(title="Trackio - Experiment Tracking", theme=gr.themes.Soft()) as
                 inputs=[hf_token_input, dataset_repo_input],
                 outputs=current_config_output
             )
+    
         
-        # Create Experiment Tab
-        with gr.Tab("Create Experiment"):
-            gr.Markdown("### Create a New Experiment")
-            with gr.Row():
-                with gr.Column():
-                    experiment_name = gr.Textbox(
-                        label="Experiment Name",
-                        placeholder="my_smollm3_finetune",
-                        value="smollm3_finetune"
-                    )
-                    experiment_description = gr.Textbox(
-                        label="Description",
-                        placeholder="Fine-tuning SmolLM3 model on custom dataset",
-                        value="SmolLM3 fine-tuning experiment"
-                    )
-                    create_btn = gr.Button("Create Experiment", variant="primary")
-                
-                with gr.Column():
-                    create_output = gr.Textbox(
-                        label="Result",
-                        lines=5,
-                        interactive=False
-                    )
-            
-            create_btn.click(
-                create_experiment_interface,
-                inputs=[experiment_name, experiment_description],
-                outputs=create_output
-            )
-        
-        # Log Metrics Tab
-        with gr.Tab("Log Metrics"):
-            gr.Markdown("### Log Training Metrics")
-            with gr.Row():
-                with gr.Column():
-                    metrics_exp_id = gr.Textbox(
-                        label="Experiment ID",
-                        placeholder="exp_20231201_143022"
-                    )
-                    metrics_json = gr.Textbox(
-                        label="Metrics (JSON)",
-                        placeholder='{"loss": 0.5, "accuracy": 0.85, "learning_rate": 2e-5}',
-                        value='{"loss": 0.5, "accuracy": 0.85, "learning_rate": 2e-5, "gpu_memory": 22.5}'
-                    )
-                    metrics_step = gr.Textbox(
-                        label="Step (optional)",
-                        placeholder="100"
-                    )
-                    log_metrics_btn = gr.Button("Log Metrics", variant="primary")
-                
-                with gr.Column():
-                    metrics_output = gr.Textbox(
-                        label="Result",
-                        lines=5,
-                        interactive=False
-                    )
-            
-            log_metrics_btn.click(
-                log_metrics_interface,
-                inputs=[metrics_exp_id, metrics_json, metrics_step],
-                outputs=metrics_output
-            )
-        
-        # Log Parameters Tab
-        with gr.Tab("Log Parameters"):
-            gr.Markdown("### Log Experiment Parameters")
-            with gr.Row():
-                with gr.Column():
-                    params_exp_id = gr.Textbox(
-                        label="Experiment ID",
-                        placeholder="exp_20231201_143022"
-                    )
-                    parameters_json = gr.Textbox(
-                        label="Parameters (JSON)",
-                        placeholder='{"learning_rate": 2e-5, "batch_size": 4}',
-                        value='{"learning_rate": 3.5e-6, "batch_size": 8, "model_name": "HuggingFaceTB/SmolLM3-3B", "max_iters": 18000, "mixed_precision": "bf16"}'
-                    )
-                    log_params_btn = gr.Button("Log Parameters", variant="primary")
-                
-                with gr.Column():
-                    params_output = gr.Textbox(
-                        label="Result",
-                        lines=5,
-                        interactive=False
-                    )
-            
-            log_params_btn.click(
-                log_parameters_interface,
-                inputs=[params_exp_id, parameters_json],
-                outputs=params_output
-            )
-        
-        # View Experiments Tab
-        with gr.Tab("View Experiments"):
-            gr.Markdown("### View Experiment Details")
-            with gr.Row():
-                with gr.Column():
-                    view_exp_id = gr.Textbox(
-                        label="Experiment ID",
-                        placeholder="exp_20231201_143022"
-                    )
-                    view_btn = gr.Button("View Experiment", variant="primary")
-                    list_btn = gr.Button("List All Experiments", variant="secondary")
-                
-                with gr.Column():
-                    view_output = gr.Textbox(
-                        label="Experiment Details",
-                        lines=20,
-                        interactive=False
-                    )
-            
-            view_btn.click(
-                get_experiment_details,
-                inputs=[view_exp_id],
-                outputs=view_output
-            )
-            
-            list_btn.click(
-                list_experiments_interface,
-                inputs=[],
-                outputs=view_output
-            )
-        
-        # Visualization Tab
-        with gr.Tab("📊 Visualizations"):
-            gr.Markdown("### Training Metrics Visualization")
-            with gr.Row():
-                with gr.Column():
-                    plot_exp_id = gr.Textbox(
-                        label="Experiment ID",
-                        placeholder="exp_20231201_143022"
-                    )
-                    metric_dropdown = gr.Dropdown(
-                        label="Metric to Plot",
-                        choices=[
-                            "loss", "accuracy", "learning_rate", "gpu_memory", "training_time",
-                            "total_tokens", "truncated_tokens", "padding_tokens", "throughput", "step_time",
-                            "batch_size", "seq_len", "token_acc", "train/gate_ortho", "train/center"
-                        ],
-                        value="loss"
-                    )
-                    plot_btn = gr.Button("Create Plot", variant="primary")
-                
-                with gr.Column():
-                    plot_output = gr.Plot(label="Training Metrics")
-            
-            plot_btn.click(
-                create_metrics_plot,
-                inputs=[plot_exp_id, metric_dropdown],
-                outputs=plot_output
-            )
-            
-            gr.Markdown("### Experiment Comparison")
-            with gr.Row():
-                with gr.Column():
-                    comparison_exp_ids = gr.Textbox(
-                        label="Experiment IDs (comma-separated)",
-                        placeholder="exp_1,exp_2,exp_3"
-                    )
-                    comparison_btn = gr.Button("Compare Experiments", variant="primary")
-                
-                with gr.Column():
-                    comparison_plot = gr.Plot(label="Experiment Comparison")
-            
-            comparison_btn.click(
-                create_experiment_comparison,
-                inputs=[comparison_exp_ids],
-                outputs=comparison_plot
-            )
-        
-        # Demo Data Tab
-        with gr.Tab("🎯 Demo Data"):
-            gr.Markdown("### Generate Demo Training Data")
-            gr.Markdown("Use this to simulate training data for testing the interface")
-            with gr.Row():
-                with gr.Column():
-                    demo_exp_id = gr.Textbox(
-                        label="Experiment ID",
-                        placeholder="exp_20231201_143022"
-                    )
-                    demo_btn = gr.Button("Generate Demo Data", variant="primary")
-                    create_demo_btn = gr.Button("Create Demo Experiment", variant="secondary")
-                
-                with gr.Column():
-                    demo_output = gr.Textbox(
-                        label="Result",
-                        lines=5,
-                        interactive=False
-                    )
-            
-            demo_btn.click(
-                simulate_training_data,
-                inputs=[demo_exp_id],
-                outputs=demo_output
-            )
-            
-            create_demo_btn.click(
-                create_demo_experiment,
-                inputs=[],
-                outputs=demo_output
-            )
-        
-        # Update Status Tab
-        with gr.Tab("Update Status"):
-            gr.Markdown("### Update Experiment Status")
-            with gr.Row():
-                with gr.Column():
-                    status_exp_id = gr.Textbox(
-                        label="Experiment ID",
-                        placeholder="exp_20231201_143022"
-                    )
-                    status_dropdown = gr.Dropdown(
-                        label="Status",
-                        choices=["running", "completed", "failed", "paused"],
-                        value="running"
-                    )
-                    update_status_btn = gr.Button("Update Status", variant="primary")
-                
-                with gr.Column():
-                    status_output = gr.Textbox(
-                        label="Result",
-                        lines=3,
-                        interactive=False
-                    )
-            
-            update_status_btn.click(
-                update_experiment_status_interface,
-                inputs=[status_exp_id, status_dropdown],
-                outputs=status_output
-            )
+
 
 # Launch the app
 if __name__ == "__main__":

templates/spaces/trackio/dataset_utils.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""
+Dataset utilities for Trackio experiment data management
+Provides functions for safe dataset operations with data preservation
+"""
+
+import json
+import logging
+from datetime import datetime
+from typing import Dict, Any, List, Optional, Union
+from datasets import Dataset, load_dataset
+
+logger = logging.getLogger(__name__)
+
+class TrackioDatasetManager:
+    """
+    Manager class for Trackio experiment datasets with data preservation.
+    
+    This class ensures that existing experiment data is always preserved
+    when adding new experiments or updating existing ones.
+    """
+    
+    def __init__(self, dataset_repo: str, hf_token: str):
+        """
+        Initialize the dataset manager.
+        
+        Args:
+            dataset_repo (str): HF dataset repository ID (e.g., "username/dataset-name")
+            hf_token (str): Hugging Face token for authentication
+        """
+        self.dataset_repo = dataset_repo
+        self.hf_token = hf_token
+        self._validate_repo_format()
+    
+    def _validate_repo_format(self):
+        """Validate dataset repository format"""
+        if not self.dataset_repo or '/' not in self.dataset_repo:
+            raise ValueError(f"Invalid dataset repository format: {self.dataset_repo}")
+    
+    def check_dataset_exists(self) -> bool:
+        """
+        Check if the dataset repository exists and is accessible.
+        
+        Returns:
+            bool: True if dataset exists and is accessible, False otherwise
+        """
+        try:
+            load_dataset(self.dataset_repo, token=self.hf_token)
+            logger.info(f"✅ Dataset {self.dataset_repo} exists and is accessible")
+            return True
+        except Exception as e:
+            logger.info(f"📊 Dataset {self.dataset_repo} doesn't exist or isn't accessible: {e}")
+            return False
+    
+    def load_existing_experiments(self) -> List[Dict[str, Any]]:
+        """
+        Load all existing experiments from the dataset.
+        
+        Returns:
+            List[Dict[str, Any]]: List of existing experiment dictionaries
+        """
+        try:
+            if not self.check_dataset_exists():
+                logger.info("📊 No existing dataset found, returning empty list")
+                return []
+            
+            dataset = load_dataset(self.dataset_repo, token=self.hf_token)
+            
+            if 'train' not in dataset:
+                logger.info("📊 No 'train' split found in dataset")
+                return []
+            
+            experiments = list(dataset['train'])
+            logger.info(f"📊 Loaded {len(experiments)} existing experiments")
+            
+            # Validate experiment structure
+            valid_experiments = []
+            for exp in experiments:
+                if self._validate_experiment_structure(exp):
+                    valid_experiments.append(exp)
+                else:
+                    logger.warning(f"⚠️ Skipping invalid experiment: {exp.get('experiment_id', 'unknown')}")
+            
+            logger.info(f"📊 {len(valid_experiments)} valid experiments loaded")
+            return valid_experiments
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to load existing experiments: {e}")
+            return []
+    
+    def _validate_experiment_structure(self, experiment: Dict[str, Any]) -> bool:
+        """
+        Validate that an experiment has the required structure.
+        
+        Args:
+            experiment (Dict[str, Any]): Experiment dictionary to validate
+            
+        Returns:
+            bool: True if experiment structure is valid
+        """
+        required_fields = [
+            'experiment_id', 'name', 'description', 'created_at', 
+            'status', 'metrics', 'parameters', 'artifacts', 'logs'
+        ]
+        
+        for field in required_fields:
+            if field not in experiment:
+                logger.warning(f"⚠️ Missing required field '{field}' in experiment")
+                return False
+        
+        # Validate JSON fields
+        json_fields = ['metrics', 'parameters', 'artifacts', 'logs']
+        for field in json_fields:
+            if isinstance(experiment[field], str):
+                try:
+                    json.loads(experiment[field])
+                except json.JSONDecodeError:
+                    logger.warning(f"⚠️ Invalid JSON in field '{field}' for experiment {experiment.get('experiment_id')}")
+                    return False
+        
+        return True
+    
+    def save_experiments(self, experiments: List[Dict[str, Any]], commit_message: Optional[str] = None) -> bool:
+        """
+        Save a list of experiments to the dataset, preserving data integrity.
+        
+        Args:
+            experiments (List[Dict[str, Any]]): List of experiment dictionaries
+            commit_message (Optional[str]): Custom commit message
+            
+        Returns:
+            bool: True if save was successful, False otherwise
+        """
+        try:
+            if not experiments:
+                logger.warning("⚠️ No experiments to save")
+                return False
+            
+            # Validate all experiments before saving
+            valid_experiments = []
+            for exp in experiments:
+                if self._validate_experiment_structure(exp):
+                    # Ensure last_updated is set
+                    if 'last_updated' not in exp:
+                        exp['last_updated'] = datetime.now().isoformat()
+                    valid_experiments.append(exp)
+                else:
+                    logger.error(f"❌ Invalid experiment structure: {exp.get('experiment_id', 'unknown')}")
+                    return False
+            
+            # Create dataset
+            dataset = Dataset.from_list(valid_experiments)
+            
+            # Generate commit message if not provided
+            if not commit_message:
+                commit_message = f"Update dataset with {len(valid_experiments)} experiments ({datetime.now().isoformat()})"
+            
+            # Push to hub
+            dataset.push_to_hub(
+                self.dataset_repo,
+                token=self.hf_token,
+                private=True,
+                commit_message=commit_message
+            )
+            
+            logger.info(f"✅ Successfully saved {len(valid_experiments)} experiments to {self.dataset_repo}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to save experiments to dataset: {e}")
+            return False
+    
+    def upsert_experiment(self, experiment: Dict[str, Any]) -> bool:
+        """
+        Insert a new experiment or update an existing one, preserving all other data.
+        
+        Args:
+            experiment (Dict[str, Any]): Experiment dictionary to upsert
+            
+        Returns:
+            bool: True if operation was successful, False otherwise
+        """
+        try:
+            # Validate the experiment structure
+            if not self._validate_experiment_structure(experiment):
+                logger.error(f"❌ Invalid experiment structure for {experiment.get('experiment_id', 'unknown')}")
+                return False
+            
+            # Load existing experiments
+            existing_experiments = self.load_existing_experiments()
+            
+            # Find if experiment already exists
+            experiment_id = experiment['experiment_id']
+            experiment_found = False
+            updated_experiments = []
+            
+            for existing_exp in existing_experiments:
+                if existing_exp.get('experiment_id') == experiment_id:
+                    # Update existing experiment
+                    logger.info(f"🔄 Updating existing experiment: {experiment_id}")
+                    experiment['last_updated'] = datetime.now().isoformat()
+                    updated_experiments.append(experiment)
+                    experiment_found = True
+                else:
+                    # Preserve existing experiment
+                    updated_experiments.append(existing_exp)
+            
+            # If experiment doesn't exist, add it
+            if not experiment_found:
+                logger.info(f"➕ Adding new experiment: {experiment_id}")
+                experiment['last_updated'] = datetime.now().isoformat()
+                updated_experiments.append(experiment)
+            
+            # Save all experiments
+            commit_message = f"{'Update' if experiment_found else 'Add'} experiment {experiment_id} (preserving {len(existing_experiments)} existing experiments)"
+            
+            return self.save_experiments(updated_experiments, commit_message)
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to upsert experiment: {e}")
+            return False
+    
+    def get_experiment_by_id(self, experiment_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve a specific experiment by its ID.
+        
+        Args:
+            experiment_id (str): The experiment ID to search for
+            
+        Returns:
+            Optional[Dict[str, Any]]: The experiment dictionary if found, None otherwise
+        """
+        try:
+            experiments = self.load_existing_experiments()
+            
+            for exp in experiments:
+                if exp.get('experiment_id') == experiment_id:
+                    logger.info(f"✅ Found experiment: {experiment_id}")
+                    return exp
+            
+            logger.info(f"📊 Experiment not found: {experiment_id}")
+            return None
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to get experiment {experiment_id}: {e}")
+            return None
+    
+    def list_experiments(self, status_filter: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        List all experiments, optionally filtered by status.
+        
+        Args:
+            status_filter (Optional[str]): Filter by experiment status (running, completed, failed, paused)
+            
+        Returns:
+            List[Dict[str, Any]]: List of experiments matching the filter
+        """
+        try:
+            experiments = self.load_existing_experiments()
+            
+            if status_filter:
+                filtered_experiments = [exp for exp in experiments if exp.get('status') == status_filter]
+                logger.info(f"📊 Found {len(filtered_experiments)} experiments with status '{status_filter}'")
+                return filtered_experiments
+            
+            logger.info(f"📊 Found {len(experiments)} total experiments")
+            return experiments
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to list experiments: {e}")
+            return []
+    
+    def backup_dataset(self, backup_suffix: Optional[str] = None) -> str:
+        """
+        Create a backup of the current dataset.
+        
+        Args:
+            backup_suffix (Optional[str]): Optional suffix for backup repo name
+            
+        Returns:
+            str: Backup repository name if successful, empty string otherwise
+        """
+        try:
+            if not backup_suffix:
+                backup_suffix = datetime.now().strftime('%Y%m%d_%H%M%S')
+            
+            backup_repo = f"{self.dataset_repo}-backup-{backup_suffix}"
+            
+            # Load current experiments
+            experiments = self.load_existing_experiments()
+            
+            if not experiments:
+                logger.warning("⚠️ No experiments to backup")
+                return ""
+            
+            # Create backup dataset manager
+            backup_manager = TrackioDatasetManager(backup_repo, self.hf_token)
+            
+            # Save to backup
+            success = backup_manager.save_experiments(
+                experiments, 
+                f"Backup of {self.dataset_repo} created on {datetime.now().isoformat()}"
+            )
+            
+            if success:
+                logger.info(f"✅ Backup created: {backup_repo}")
+                return backup_repo
+            else:
+                logger.error("❌ Failed to create backup")
+                return ""
+                
+        except Exception as e:
+            logger.error(f"❌ Failed to create backup: {e}")
+            return ""
+
+
+def create_dataset_manager(dataset_repo: str, hf_token: str) -> TrackioDatasetManager:
+    """
+    Factory function to create a TrackioDatasetManager instance.
+    
+    Args:
+        dataset_repo (str): HF dataset repository ID
+        hf_token (str): Hugging Face token
+        
+    Returns:
+        TrackioDatasetManager: Configured dataset manager instance
+    """
+    return TrackioDatasetManager(dataset_repo, hf_token)

templates/spaces/trackio/trackio_api_client.py

@@ -0,0 +1,320 @@
+#!/usr/bin/env python3
+"""
+Trackio API Client for Hugging Face Spaces
+Uses gradio_client for proper API communication with automatic Space URL resolution
+"""
+
+import requests
+import json
+import time
+import logging
+from typing import Dict, Any, Optional
+from datetime import datetime
+import os
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+try:
+    from gradio_client import Client
+    GRADIO_CLIENT_AVAILABLE = True
+except ImportError:
+    GRADIO_CLIENT_AVAILABLE = False
+    logger.warning("gradio_client not available. Install with: pip install gradio_client")
+
+try:
+    from huggingface_hub import HfApi
+    HF_HUB_AVAILABLE = True
+except ImportError:
+    HF_HUB_AVAILABLE = False
+    logger.warning("huggingface_hub not available. Install with: pip install huggingface-hub")
+
+class TrackioAPIClient:
+    """API client for Trackio Space using gradio_client with automatic Space URL resolution"""
+    
+    def __init__(self, space_id: str, hf_token: Optional[str] = None):
+        self.space_id = space_id
+        self.hf_token = hf_token
+        self.client = None
+        
+        # Auto-resolve Space URL
+        self.space_url = self._resolve_space_url()
+        
+        # Initialize gradio client
+        if GRADIO_CLIENT_AVAILABLE and self.space_url:
+            try:
+                self.client = Client(self.space_url)
+                logger.info(f"✅ Connected to Trackio Space: {self.space_id}")
+            except Exception as e:
+                logger.error(f"❌ Failed to connect to Trackio Space: {e}")
+                self.client = None
+        else:
+            logger.error("❌ gradio_client not available. Install with: pip install gradio_client")
+    
+    def _resolve_space_url(self) -> Optional[str]:
+        """Resolve Space URL using Hugging Face Hub API"""
+        try:
+            # Clean the space_id - remove any URL prefixes
+            clean_space_id = self.space_id
+            if clean_space_id.startswith('http'):
+                # Extract space ID from URL
+                if '/spaces/' in clean_space_id:
+                    clean_space_id = clean_space_id.split('/spaces/')[-1]
+                else:
+                    # Try to extract from URL format
+                    clean_space_id = clean_space_id.replace('https://', '').replace('http://', '')
+                    if '.hf.space' in clean_space_id:
+                        clean_space_id = clean_space_id.replace('.hf.space', '').replace('-', '/')
+            
+            logger.info(f"🔧 Resolving Space URL for ID: {clean_space_id}")
+            
+            if not HF_HUB_AVAILABLE:
+                logger.warning("⚠️ Hugging Face Hub not available, using default URL format")
+                # Fallback to default URL format
+                space_name = clean_space_id.replace('/', '-')
+                return f"https://{space_name}.hf.space"
+            
+            # Use Hugging Face Hub API to get Space info
+            api = HfApi(token=self.hf_token)
+            
+            # Get Space info
+            space_info = api.space_info(clean_space_id)
+            if space_info and hasattr(space_info, 'host'):
+                # Use the host directly from space_info
+                space_url = space_info.host
+                logger.info(f"✅ Resolved Space URL: {space_url}")
+                return space_url
+            else:
+                # Fallback to default URL format
+                space_name = clean_space_id.replace('/', '-')
+                space_url = f"https://{space_name}.hf.space"
+                logger.info(f"✅ Using fallback Space URL: {space_url}")
+                return space_url
+                
+        except Exception as e:
+            logger.warning(f"⚠️ Failed to resolve Space URL: {e}")
+            # Fallback to default URL format
+            space_name = self.space_id.replace('/', '-')
+            space_url = f"https://{space_name}.hf.space"
+            logger.info(f"✅ Using fallback Space URL: {space_url}")
+            return space_url
+    
+    def _make_api_call(self, api_name: str, *args) -> Dict[str, Any]:
+        """Make an API call to the Trackio Space using gradio_client"""
+        if not self.client:
+            return {"error": "Client not available"}
+        
+        try:
+            logger.debug(f"Making API call to {api_name} with args: {args}")
+            
+            # Use gradio_client to make the prediction
+            result = self.client.predict(*args, api_name=api_name)
+            
+            logger.debug(f"API call result: {result}")
+            return {"success": True, "data": result}
+            
+        except Exception as e:
+            logger.error(f"API call failed for {api_name}: {e}")
+            return {"error": f"API call failed: {str(e)}"}
+    
+    def create_experiment(self, name: str, description: str = "") -> Dict[str, Any]:
+        """Create a new experiment"""
+        logger.info(f"Creating experiment: {name}")
+        
+        result = self._make_api_call("/create_experiment_interface", name, description)
+        
+        if "success" in result:
+            logger.info(f"Experiment created successfully: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to create experiment: {result}")
+            return result
+    
+    def log_metrics(self, experiment_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> Dict[str, Any]:
+        """Log metrics for an experiment"""
+        metrics_json = json.dumps(metrics)
+        step_str = str(step) if step is not None else ""
+        
+        logger.info(f"Logging metrics for experiment {experiment_id} at step {step}")
+        
+        result = self._make_api_call("/log_metrics_interface", experiment_id, metrics_json, step_str)
+        
+        if "success" in result:
+            logger.info(f"Metrics logged successfully: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to log metrics: {result}")
+            return result
+    
+    def log_parameters(self, experiment_id: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
+        """Log parameters for an experiment"""
+        parameters_json = json.dumps(parameters)
+        
+        logger.info(f"Logging parameters for experiment {experiment_id}")
+        
+        result = self._make_api_call("/log_parameters_interface", experiment_id, parameters_json)
+        
+        if "success" in result:
+            logger.info(f"Parameters logged successfully: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to log parameters: {result}")
+            return result
+    
+    def get_experiment_details(self, experiment_id: str) -> Dict[str, Any]:
+        """Get experiment details"""
+        logger.info(f"Getting details for experiment {experiment_id}")
+        
+        result = self._make_api_call("/get_experiment_details", experiment_id)
+        
+        if "success" in result:
+            logger.info(f"Experiment details retrieved: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to get experiment details: {result}")
+            return result
+    
+    def list_experiments(self) -> Dict[str, Any]:
+        """List all experiments"""
+        logger.info("Listing experiments")
+        
+        result = self._make_api_call("/list_experiments_interface")
+        
+        if "success" in result:
+            logger.info(f"Experiments listed successfully: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to list experiments: {result}")
+            return result
+    
+    def update_experiment_status(self, experiment_id: str, status: str) -> Dict[str, Any]:
+        """Update experiment status"""
+        logger.info(f"Updating experiment {experiment_id} status to {status}")
+        
+        result = self._make_api_call("/update_experiment_status_interface", experiment_id, status)
+        
+        if "success" in result:
+            logger.info(f"Experiment status updated successfully: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to update experiment status: {result}")
+            return result
+    
+    def simulate_training_data(self, experiment_id: str) -> Dict[str, Any]:
+        """Simulate training data for testing"""
+        logger.info(f"Simulating training data for experiment {experiment_id}")
+        
+        result = self._make_api_call("/simulate_training_data", experiment_id)
+        
+        if "success" in result:
+            logger.info(f"Training data simulated successfully: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to simulate training data: {result}")
+            return result
+    
+    def get_training_metrics(self, experiment_id: str) -> Dict[str, Any]:
+        """Get training metrics for an experiment"""
+        logger.info(f"Getting training metrics for experiment {experiment_id}")
+        
+        result = self._make_api_call("/get_experiment_details", experiment_id)
+        
+        if "success" in result:
+            logger.info(f"Training metrics retrieved: {result['data']}")
+            return result
+        else:
+            logger.error(f"Failed to get training metrics: {result}")
+            return result
+    
+    def create_metrics_plot(self, experiment_id: str, metric_name: str = "loss") -> Dict[str, Any]:
+        """Create a metrics plot for an experiment"""
+        logger.info(f"Creating metrics plot for experiment {experiment_id}, metric: {metric_name}")
+        
+        result = self._make_api_call("/create_metrics_plot", experiment_id, metric_name)
+        
+        if "success" in result:
+            logger.info(f"Metrics plot created successfully")
+            return result
+        else:
+            logger.error(f"Failed to create metrics plot: {result}")
+            return result
+    
+    def create_experiment_comparison(self, experiment_ids: str) -> Dict[str, Any]:
+        """Compare multiple experiments"""
+        logger.info(f"Creating experiment comparison for: {experiment_ids}")
+        
+        result = self._make_api_call("/create_experiment_comparison", experiment_ids)
+        
+        if "success" in result:
+            logger.info(f"Experiment comparison created successfully")
+            return result
+        else:
+            logger.error(f"Failed to create experiment comparison: {result}")
+            return result
+    
+    def test_connection(self) -> Dict[str, Any]:
+        """Test connection to the Trackio Space"""
+        logger.info("Testing connection to Trackio Space")
+        
+        try:
+            # Try to list experiments as a connection test
+            result = self.list_experiments()
+            if "success" in result:
+                return {"success": True, "message": "Connection successful"}
+            else:
+                return {"error": "Connection failed", "details": result}
+        except Exception as e:
+            return {"error": f"Connection test failed: {str(e)}"}
+    
+    def get_space_info(self) -> Dict[str, Any]:
+        """Get information about the Space"""
+        try:
+            if not HF_HUB_AVAILABLE:
+                return {"error": "Hugging Face Hub not available"}
+            
+            api = HfApi(token=self.hf_token)
+            space_info = api.space_info(self.space_id)
+            
+            return {
+                "success": True,
+                "data": {
+                    "space_id": self.space_id,
+                    "space_url": self.space_url,
+                    "space_info": {
+                        "title": getattr(space_info, 'title', 'Unknown'),
+                        "host": getattr(space_info, 'host', 'Unknown'),
+                        "stage": getattr(space_info, 'stage', 'Unknown'),
+                        "visibility": getattr(space_info, 'visibility', 'Unknown')
+                    }
+                }
+            }
+        except Exception as e:
+            return {"error": f"Failed to get Space info: {str(e)}"}
+
+# Factory function to create client with dynamic configuration
+def create_trackio_client(space_id: Optional[str] = None, hf_token: Optional[str] = None) -> TrackioAPIClient:
+    """Create a TrackioAPIClient with dynamic configuration"""
+    
+    # Get space_id from environment if not provided
+    if not space_id:
+        space_id = os.environ.get('TRACKIO_URL')
+        if not space_id:
+            # Try to construct from username and space name
+            username = os.environ.get('HF_USERNAME')
+            space_name = os.environ.get('TRACKIO_SPACE_NAME')
+            if username and space_name:
+                space_id = f"https://huggingface.co/spaces/{username}/{space_name}"
+            else:
+                logger.warning("⚠️ No space_id provided and could not determine from environment")
+                return None
+    
+    # Get HF token from environment if not provided
+    if not hf_token:
+        hf_token = os.environ.get('HF_TOKEN')
+    
+    if not space_id:
+        logger.error("❌ No space_id available for TrackioAPIClient")
+        return None
+    
+    return TrackioAPIClient(space_id, hf_token) 

tests/test_data_preservation.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+"""
+Test script to validate data preservation in Trackio dataset operations
+"""
+
+import os
+import sys
+import json
+import tempfile
+import logging
+from datetime import datetime
+from typing import Dict, Any
+
+# Add src to path for imports
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+
+from dataset_utils import TrackioDatasetManager
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def create_sample_experiment(experiment_id: str, name: str, status: str = "running") -> Dict[str, Any]:
+    """Create a sample experiment for testing"""
+    return {
+        'experiment_id': experiment_id,
+        'name': name,
+        'description': f"Test experiment {name}",
+        'created_at': datetime.now().isoformat(),
+        'status': status,
+        'metrics': json.dumps([
+            {
+                'timestamp': datetime.now().isoformat(),
+                'step': 100,
+                'metrics': {
+                    'loss': 1.5,
+                    'accuracy': 0.85,
+                    'learning_rate': 5e-6
+                }
+            }
+        ]),
+        'parameters': json.dumps({
+            'model_name': 'HuggingFaceTB/SmolLM3-3B',
+            'batch_size': 8,
+            'learning_rate': 5e-6
+        }),
+        'artifacts': json.dumps([]),
+        'logs': json.dumps([]),
+        'last_updated': datetime.now().isoformat()
+    }
+
+def test_data_preservation():
+    """Test data preservation functionality"""
+    # Get HF token from environment
+    hf_token = os.environ.get('HF_TOKEN') or os.environ.get('HUGGING_FACE_HUB_TOKEN')
+    
+    if not hf_token:
+        logger.error("❌ HF_TOKEN not found in environment variables")
+        logger.info("Please set HF_TOKEN or HUGGING_FACE_HUB_TOKEN environment variable")
+        return False
+    
+    # Use a test dataset repository
+    test_dataset_repo = "tonic/trackio-test-preservation"
+    
+    try:
+        logger.info("🧪 Starting data preservation test")
+        logger.info(f"📊 Test dataset: {test_dataset_repo}")
+        
+        # Initialize dataset manager
+        dataset_manager = TrackioDatasetManager(test_dataset_repo, hf_token)
+        
+        # Test 1: Check if dataset exists
+        logger.info("\n📝 Test 1: Checking dataset existence...")
+        exists = dataset_manager.check_dataset_exists()
+        logger.info(f"Dataset exists: {exists}")
+        
+        # Test 2: Load existing experiments (should handle empty/non-existent gracefully)
+        logger.info("\n📝 Test 2: Loading existing experiments...")
+        existing_experiments = dataset_manager.load_existing_experiments()
+        logger.info(f"Found {len(existing_experiments)} existing experiments")
+        
+        # Test 3: Add first experiment
+        logger.info("\n📝 Test 3: Adding first experiment...")
+        exp1 = create_sample_experiment("test_exp_001", "First Test Experiment")
+        success = dataset_manager.upsert_experiment(exp1)
+        logger.info(f"First experiment added: {success}")
+        
+        if not success:
+            logger.error("❌ Failed to add first experiment")
+            return False
+        
+        # Test 4: Add second experiment (should preserve first)
+        logger.info("\n📝 Test 4: Adding second experiment...")
+        exp2 = create_sample_experiment("test_exp_002", "Second Test Experiment")
+        success = dataset_manager.upsert_experiment(exp2)
+        logger.info(f"Second experiment added: {success}")
+        
+        if not success:
+            logger.error("❌ Failed to add second experiment")
+            return False
+        
+        # Test 5: Verify both experiments exist
+        logger.info("\n📝 Test 5: Verifying both experiments exist...")
+        all_experiments = dataset_manager.load_existing_experiments()
+        logger.info(f"Total experiments after adding two: {len(all_experiments)}")
+        
+        exp_ids = [exp.get('experiment_id') for exp in all_experiments]
+        if "test_exp_001" in exp_ids and "test_exp_002" in exp_ids:
+            logger.info("✅ Both experiments preserved successfully")
+        else:
+            logger.error(f"❌ Experiments not preserved. Found IDs: {exp_ids}")
+            return False
+        
+        # Test 6: Update existing experiment (should preserve others)
+        logger.info("\n📝 Test 6: Updating first experiment...")
+        exp1_updated = create_sample_experiment("test_exp_001", "Updated First Experiment", "completed")
+        success = dataset_manager.upsert_experiment(exp1_updated)
+        logger.info(f"First experiment updated: {success}")
+        
+        if not success:
+            logger.error("❌ Failed to update first experiment")
+            return False
+        
+        # Test 7: Verify update preserved other experiments
+        logger.info("\n📝 Test 7: Verifying update preserved other experiments...")
+        final_experiments = dataset_manager.load_existing_experiments()
+        logger.info(f"Total experiments after update: {len(final_experiments)}")
+        
+        # Check that we still have both experiments
+        if len(final_experiments) != 2:
+            logger.error(f"❌ Wrong number of experiments after update: {len(final_experiments)}")
+            return False
+        
+        # Check that first experiment was updated
+        exp1_final = dataset_manager.get_experiment_by_id("test_exp_001")
+        if exp1_final and exp1_final.get('status') == 'completed':
+            logger.info("✅ First experiment successfully updated")
+        else:
+            logger.error("❌ First experiment update failed")
+            return False
+        
+        # Check that second experiment was preserved
+        exp2_final = dataset_manager.get_experiment_by_id("test_exp_002")
+        if exp2_final and exp2_final.get('name') == "Second Test Experiment":
+            logger.info("✅ Second experiment successfully preserved")
+        else:
+            logger.error("❌ Second experiment not preserved")
+            return False
+        
+        # Test 8: Test filtering functionality
+        logger.info("\n📝 Test 8: Testing filtering functionality...")
+        running_experiments = dataset_manager.list_experiments(status_filter="running")
+        completed_experiments = dataset_manager.list_experiments(status_filter="completed")
+        
+        logger.info(f"Running experiments: {len(running_experiments)}")
+        logger.info(f"Completed experiments: {len(completed_experiments)}")
+        
+        if len(running_experiments) == 1 and len(completed_experiments) == 1:
+            logger.info("✅ Filtering functionality works correctly")
+        else:
+            logger.error("❌ Filtering functionality failed")
+            return False
+        
+        logger.info("\n🎉 All data preservation tests passed!")
+        logger.info("✅ Data preservation functionality is working correctly")
+        return True
+        
+    except Exception as e:
+        logger.error(f"❌ Test failed with exception: {e}")
+        return False
+
+def main():
+    """Main test function"""
+    logger.info("Data Preservation Test Suite")
+    logger.info("=" * 50)
+    
+    success = test_data_preservation()
+    
+    if success:
+        logger.info("\n✅ All tests passed!")
+        sys.exit(0)
+    else:
+        logger.error("\n❌ Some tests failed!")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()

tests/test_demo_deployment.py

@@ -39,14 +39,17 @@ def test_template_files_exist():
     """Test that template files exist"""
     print("🧪 Testing template files existence...")
     
-    template_dir = Path(__file__).parent.parent / "templates" / "spaces" / "demo"
-    
+    demo_types = ["demo_smol", "demo_gpt"]
     required_files = ["app.py", "requirements.txt"]
     
-    for file_name in required_files:
-        file_path = template_dir / file_name
-        assert file_path.exists(), f"Required file {file_name} not found in templates"
-        print(f"✅ Found {file_name}")
+    for demo_type in demo_types:
+        template_dir = Path(__file__).parent.parent / "templates" / "spaces" / demo_type
+        print(f"Checking {demo_type} templates...")
+        
+        for file_name in required_files:
+            file_path = template_dir / file_name
+            assert file_path.exists(), f"Required file {file_name} not found in {demo_type} templates"
+            print(f"✅ Found {demo_type}/{file_name}")
     
     print("✅ Template files test passed")
 

tests/test_deployment.py

@@ -17,15 +17,19 @@ def test_templates_exist():
     
     # Check spaces templates
     spaces_dir = project_root / "templates" / "spaces"
+    demo_types = ["demo_smol", "demo_gpt", "trackio"]
     spaces_files = ["app.py", "requirements.txt", "README.md"]
     
-    for file_name in spaces_files:
-        file_path = spaces_dir / file_name
-        if file_path.exists():
-            print(f"✅ {file_path}")
-        else:
-            print(f"❌ {file_path} not found")
-            return False
+    for demo_type in demo_types:
+        demo_dir = spaces_dir / demo_type
+        print(f"Checking {demo_type} templates...")
+        for file_name in spaces_files:
+            file_path = demo_dir / file_name
+            if file_path.exists():
+                print(f"✅ {file_path}")
+            else:
+                print(f"❌ {file_path} not found")
+                return False
     
     # Check datasets templates
     datasets_dir = project_root / "templates" / "datasets"

tests/test_hf_datasets.py

@@ -76,7 +76,7 @@ def test_backup_fallback():
     
     try:
         # Import and test the TrackioSpace class
-        from templates.spaces.app import TrackioSpace
+        from templates.spaces.trackio.app import TrackioSpace
         
         trackio = TrackioSpace()
         experiments = trackio.experiments
@@ -105,7 +105,7 @@ def test_metrics_dataframe():
     print("=" * 40)
     
     try:
-        from templates.spaces.app import TrackioSpace
+        from templates.spaces.trackio.app import TrackioSpace
         
         trackio = TrackioSpace()
         

tests/test_latest_deployment.py

@@ -158,16 +158,20 @@ def test_template_files():
     """Test that all required template files exist"""
     print("\n🔍 Testing template files...")
     
-    templates_dir = project_root / "templates" / "spaces"
+    spaces_dir = project_root / "templates" / "spaces"
+    demo_types = ["demo_smol", "demo_gpt", "trackio"]
     required_files = ["app.py", "requirements.txt", "README.md"]
     
-    for file_name in required_files:
-        file_path = templates_dir / file_name
-        if file_path.exists():
-            print(f"✅ {file_name} exists")
-        else:
-            print(f"❌ {file_name} missing")
-            return False
+    for demo_type in demo_types:
+        demo_dir = spaces_dir / demo_type
+        print(f"Checking {demo_type} templates...")
+        for file_name in required_files:
+            file_path = demo_dir / file_name
+            if file_path.exists():
+                print(f"✅ {demo_type}/{file_name} exists")
+            else:
+                print(f"❌ {demo_type}/{file_name} missing")
+                return False
     
     return True
 

tests/test_readme_template.py

@@ -16,8 +16,8 @@ def test_readme_template():
     print("🔍 Testing README template replacement...")
     
     try:
-        # Get template path
-        templates_dir = project_root / "templates" / "spaces"
+        # Get template path (using trackio as example)
+        templates_dir = project_root / "templates" / "spaces" / "trackio"
         readme_template_path = templates_dir / "README.md"
         
         if not readme_template_path.exists():

tests/test_real_dataset_access.py

@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+"""
+Test script to verify that the Trackio Space can read from the real Hugging Face dataset
+This test requires an HF_TOKEN environment variable to access the dataset
+"""
+
+import sys
+import os
+import json
+import logging
+from typing import Dict, Any
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def test_direct_dataset_access():
+    """Test direct access to the Hugging Face dataset"""
+    try:
+        hf_token = os.environ.get('HF_TOKEN')
+        
+        if not hf_token:
+            logger.warning("⚠️ No HF_TOKEN found. Skipping real dataset test.")
+            logger.info("💡 Set HF_TOKEN environment variable to test with real dataset")
+            return False
+        
+        from datasets import load_dataset
+        
+        dataset_repo = "Tonic/trackio-experiments"
+        logger.info(f"🔧 Testing direct access to {dataset_repo}")
+        
+        # Load the dataset
+        dataset = load_dataset(dataset_repo, token=hf_token)
+        
+        # Check structure
+        experiment_count = len(dataset['train']) if 'train' in dataset else 0
+        logger.info(f"📊 Dataset contains {experiment_count} experiments")
+        
+        if experiment_count == 0:
+            logger.warning("⚠️ No experiments found in dataset")
+            return False
+        
+        # Check columns
+        columns = list(dataset['train'].column_names) if 'train' in dataset else []
+        logger.info(f"📋 Dataset columns: {columns}")
+        
+        expected_columns = ['experiment_id', 'name', 'description', 'created_at', 'status', 'metrics', 'parameters', 'artifacts', 'logs', 'last_updated']
+        missing_columns = [col for col in expected_columns if col not in columns]
+        
+        if missing_columns:
+            logger.warning(f"⚠️ Missing expected columns: {missing_columns}")
+        else:
+            logger.info("✅ All expected columns present")
+        
+        # Test parsing a few experiments
+        successful_parses = 0
+        for i, row in enumerate(dataset['train']):
+            if i >= 3:  # Test first 3 experiments
+                break
+                
+            exp_id = row.get('experiment_id', 'unknown')
+            logger.info(f"\n🔬 Testing experiment: {exp_id}")
+            
+            # Test metrics parsing
+            metrics_raw = row.get('metrics', '[]')
+            try:
+                if isinstance(metrics_raw, str):
+                    metrics = json.loads(metrics_raw)
+                    if isinstance(metrics, list):
+                        logger.info(f"   ✅ Metrics parsed: {len(metrics)} entries")
+                        if metrics:
+                            first_metric = metrics[0]
+                            if 'metrics' in first_metric:
+                                metric_keys = list(first_metric['metrics'].keys())
+                                logger.info(f"   📊 Sample metrics: {metric_keys[:5]}...")
+                        successful_parses += 1
+                    else:
+                        logger.warning(f"   ⚠️ Metrics is not a list: {type(metrics)}")
+                else:
+                    logger.warning(f"   ⚠️ Metrics is not a string: {type(metrics_raw)}")
+            except json.JSONDecodeError as e:
+                logger.warning(f"   ❌ Failed to parse metrics JSON: {e}")
+            
+            # Test parameters parsing
+            parameters_raw = row.get('parameters', '{}')
+            try:
+                if isinstance(parameters_raw, str):
+                    parameters = json.loads(parameters_raw)
+                    if isinstance(parameters, dict):
+                        logger.info(f"   ✅ Parameters parsed: {len(parameters)} entries")
+                    else:
+                        logger.warning(f"   ⚠️ Parameters is not a dict: {type(parameters)}")
+                else:
+                    logger.warning(f"   ⚠️ Parameters is not a string: {type(parameters_raw)}")
+            except json.JSONDecodeError as e:
+                logger.warning(f"   ❌ Failed to parse parameters JSON: {e}")
+        
+        logger.info(f"\n📋 Successfully parsed {successful_parses} out of {min(3, experiment_count)} test experiments")
+        
+        return successful_parses > 0
+        
+    except Exception as e:
+        logger.error(f"❌ Error testing direct dataset access: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+def test_trackio_space_with_real_dataset():
+    """Test TrackioSpace class with real dataset"""
+    try:
+        hf_token = os.environ.get('HF_TOKEN')
+        
+        if not hf_token:
+            logger.warning("⚠️ No HF_TOKEN found. Skipping TrackioSpace test with real dataset.")
+            return False
+        
+        # Add the templates/spaces/trackio directory to the path
+        sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'templates', 'spaces', 'trackio'))
+        
+        from app import TrackioSpace
+        
+        dataset_repo = "Tonic/trackio-experiments"
+        logger.info(f"🔧 Testing TrackioSpace with {dataset_repo}")
+        
+        # Create TrackioSpace instance with real credentials
+        trackio_space = TrackioSpace(hf_token=hf_token, dataset_repo=dataset_repo)
+        
+        # Check if it loaded experiments from the dataset (not backup)
+        experiments_count = len(trackio_space.experiments)
+        logger.info(f"📊 TrackioSpace loaded {experiments_count} experiments")
+        
+        if experiments_count == 0:
+            logger.warning("⚠️ TrackioSpace loaded no experiments")
+            return False
+        
+        # Check if the dataset manager is available
+        if trackio_space.dataset_manager:
+            logger.info("✅ Dataset manager is available - data preservation enabled")
+        else:
+            logger.warning("⚠️ Dataset manager not available - using legacy mode")
+        
+        # Test loading a specific experiment
+        experiment_ids = list(trackio_space.experiments.keys())
+        if experiment_ids:
+            test_exp_id = experiment_ids[0]
+            logger.info(f"🔬 Testing metrics loading for {test_exp_id}")
+            
+            from app import get_metrics_dataframe
+            df = get_metrics_dataframe(test_exp_id)
+            
+            if not df.empty:
+                logger.info(f"✅ Metrics DataFrame created: {len(df)} rows, {len(df.columns)} columns")
+                logger.info(f"📊 Available metrics: {list(df.columns)}")
+                return True
+            else:
+                logger.warning(f"⚠️ Metrics DataFrame is empty for {test_exp_id}")
+                return False
+        else:
+            logger.warning("⚠️ No experiments available for testing")
+            return False
+        
+    except Exception as e:
+        logger.error(f"❌ Error testing TrackioSpace with real dataset: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+if __name__ == "__main__":
+    logger.info("🚀 Starting real dataset access test")
+    
+    # Test direct dataset access
+    logger.info("\n" + "="*60)
+    logger.info("TEST 1: Direct Dataset Access")
+    logger.info("="*60)
+    
+    direct_test_passed = test_direct_dataset_access()
+    
+    # Test TrackioSpace with real dataset
+    logger.info("\n" + "="*60)
+    logger.info("TEST 2: TrackioSpace with Real Dataset")
+    logger.info("="*60)
+    
+    trackio_test_passed = test_trackio_space_with_real_dataset()
+    
+    # Summary
+    logger.info("\n" + "="*60)
+    logger.info("TEST SUMMARY")
+    logger.info("="*60)
+    
+    logger.info(f"Direct Dataset Access: {'✅ PASSED' if direct_test_passed else '❌ FAILED/SKIPPED'}")
+    logger.info(f"TrackioSpace Integration: {'✅ PASSED' if trackio_test_passed else '❌ FAILED/SKIPPED'}")
+    
+    if direct_test_passed and trackio_test_passed:
+        logger.info("🎉 All tests passed! The dataset integration is working correctly.")
+        sys.exit(0)
+    elif not os.environ.get('HF_TOKEN'):
+        logger.info("ℹ️ Tests skipped due to missing HF_TOKEN. Set the token to test with real dataset.")
+        sys.exit(0)
+    else:
+        logger.error("❌ Some tests failed. Please check the implementation.")
+        sys.exit(1)

tests/test_trackio_dataset_fix.py

@@ -0,0 +1,167 @@
+#!/usr/bin/env python3
+"""
+Test script to verify that the Trackio Space can properly read from the actual dataset
+"""
+
+import sys
+import os
+import json
+import logging
+from typing import Dict, Any
+
+# Add the templates/spaces/trackio directory to the path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'templates', 'spaces', 'trackio'))
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def test_dataset_loading():
+    """Test loading experiments from the actual dataset"""
+    try:
+        # Import the TrackioSpace class
+        from app import TrackioSpace
+        
+        # Create a TrackioSpace instance pointing to the real dataset
+        dataset_repo = "Tonic/trackio-experiments"
+        hf_token = os.environ.get('HF_TOKEN')
+        
+        if not hf_token:
+            logger.warning("⚠️ No HF_TOKEN found in environment. Testing with public access.")
+        
+        logger.info(f"🔧 Testing dataset loading from {dataset_repo}")
+        
+        # Create TrackioSpace instance
+        trackio_space = TrackioSpace(hf_token=hf_token, dataset_repo=dataset_repo)
+        
+        # Check how many experiments were loaded
+        experiments_count = len(trackio_space.experiments)
+        logger.info(f"📊 Loaded {experiments_count} experiments")
+        
+        if experiments_count == 0:
+            logger.warning("⚠️ No experiments loaded - this might indicate a problem")
+            return False
+        
+        # Test specific experiment IDs from the logs
+        test_experiment_ids = [
+            'exp_20250720_130853',
+            'exp_20250720_134319',
+            'exp_20250727_172507',
+            'exp_20250727_172526'
+        ]
+        
+        found_experiments = []
+        for exp_id in test_experiment_ids:
+            if exp_id in trackio_space.experiments:
+                found_experiments.append(exp_id)
+                experiment = trackio_space.experiments[exp_id]
+                
+                logger.info(f"✅ Found experiment: {exp_id}")
+                logger.info(f"   Name: {experiment.get('name', 'N/A')}")
+                logger.info(f"   Status: {experiment.get('status', 'N/A')}")
+                logger.info(f"   Metrics count: {len(experiment.get('metrics', []))}")
+                logger.info(f"   Parameters count: {len(experiment.get('parameters', {}))}")
+                
+                # Test metrics parsing specifically
+                metrics = experiment.get('metrics', [])
+                if metrics:
+                    logger.info(f"   First metric entry: {metrics[0] if metrics else 'None'}")
+                    
+                    # Test if we can get a DataFrame for this experiment
+                    from app import get_metrics_dataframe
+                    df = get_metrics_dataframe(exp_id)
+                    if not df.empty:
+                        logger.info(f"   ✅ DataFrame created successfully: {len(df)} rows, {len(df.columns)} columns")
+                        logger.info(f"   Available metrics: {list(df.columns)}")
+                    else:
+                        logger.warning(f"   ⚠️ DataFrame is empty for {exp_id}")
+                else:
+                    logger.warning(f"   ⚠️ No metrics found for {exp_id}")
+        
+        logger.info(f"📋 Found {len(found_experiments)} out of {len(test_experiment_ids)} test experiments")
+        
+        if found_experiments:
+            logger.info("✅ Dataset loading appears to be working correctly!")
+            return True
+        else:
+            logger.warning("⚠️ No test experiments found - dataset loading may have issues")
+            return False
+            
+    except Exception as e:
+        logger.error(f"❌ Error testing dataset loading: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+def test_metrics_parsing():
+    """Test parsing metrics from the actual dataset format"""
+    try:
+        # Test with actual data structure from the dataset
+        sample_metrics_json = '''[{"timestamp": "2025-07-20T11:20:01.780908", "step": 25, "metrics": {"loss": 1.1659, "grad_norm": 10.3125, "learning_rate": 7e-08, "num_tokens": 1642080.0, "mean_token_accuracy": 0.75923578992486, "epoch": 0.004851130919895701}}, {"timestamp": "2025-07-20T11:26:39.042155", "step": 50, "metrics": {"loss": 1.165, "grad_norm": 10.75, "learning_rate": 1.4291666666666667e-07, "num_tokens": 3324682.0, "mean_token_accuracy": 0.7577659255266189, "epoch": 0.009702261839791402}}]'''
+        
+        logger.info("🔧 Testing metrics parsing")
+        
+        # Parse the JSON
+        metrics_list = json.loads(sample_metrics_json)
+        logger.info(f"📊 Parsed {len(metrics_list)} metric entries")
+        
+        # Convert to DataFrame format (like the app does)
+        import pandas as pd
+        df_data = []
+        for metric_entry in metrics_list:
+            if isinstance(metric_entry, dict):
+                step = metric_entry.get('step', 0)
+                timestamp = metric_entry.get('timestamp', '')
+                metrics = metric_entry.get('metrics', {})
+                
+                row = {'step': step, 'timestamp': timestamp}
+                row.update(metrics)
+                df_data.append(row)
+        
+        if df_data:
+            df = pd.DataFrame(df_data)
+            logger.info(f"✅ DataFrame created: {len(df)} rows, {len(df.columns)} columns")
+            logger.info(f"📋 Columns: {list(df.columns)}")
+            logger.info(f"📊 Sample data:\n{df.head()}")
+            return True
+        else:
+            logger.warning("⚠️ No data converted to DataFrame format")
+            return False
+            
+    except Exception as e:
+        logger.error(f"❌ Error testing metrics parsing: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+if __name__ == "__main__":
+    logger.info("🚀 Starting Trackio dataset fix verification")
+    
+    # Test metrics parsing first
+    logger.info("\n" + "="*50)
+    logger.info("TEST 1: Metrics Parsing")
+    logger.info("="*50)
+    
+    metrics_test_passed = test_metrics_parsing()
+    
+    # Test dataset loading
+    logger.info("\n" + "="*50)
+    logger.info("TEST 2: Dataset Loading")
+    logger.info("="*50)
+    
+    dataset_test_passed = test_dataset_loading()
+    
+    # Summary
+    logger.info("\n" + "="*50)
+    logger.info("TEST SUMMARY")
+    logger.info("="*50)
+    
+    logger.info(f"Metrics Parsing: {'✅ PASSED' if metrics_test_passed else '❌ FAILED'}")
+    logger.info(f"Dataset Loading: {'✅ PASSED' if dataset_test_passed else '❌ FAILED'}")
+    
+    if metrics_test_passed and dataset_test_passed:
+        logger.info("🎉 All tests passed! The dataset fix should work correctly.")
+        sys.exit(0)
+    else:
+        logger.error("❌ Some tests failed. Please check the implementation.")
+        sys.exit(1)

tests/test_trackio_deployment.py

@@ -17,16 +17,16 @@ def test_templates_structure():
     """Test that the templates structure is correct"""
     print("🔍 Testing templates structure...")
     
-    templates_dir = project_root / "templates" / "spaces"
+    trackio_dir = project_root / "templates" / "spaces" / "trackio"
     
     required_files = ["app.py", "requirements.txt", "README.md"]
     
     for file_name in required_files:
-        file_path = templates_dir / file_name
+        file_path = trackio_dir / file_name
         if file_path.exists():
-            print(f"✅ {file_name} exists")
+            print(f"✅ trackio/{file_name} exists")
         else:
-            print(f"❌ {file_name} missing")
+            print(f"❌ trackio/{file_name} missing")
             return False
     
     return True
@@ -35,7 +35,7 @@ def test_app_py_content():
     """Test that app.py has the required structure"""
     print("\n🔍 Testing app.py content...")
     
-    app_path = project_root / "templates" / "spaces" / "app.py"
+    app_path = project_root / "templates" / "spaces" / "trackio" / "app.py"
     
     try:
         with open(app_path, 'r', encoding='utf-8') as f:

tests/test_trackio_space_diagnostics.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+"""
+Diagnostic script for Trackio Space issues
+Helps debug dataset loading and API client issues
+"""
+
+import os
+import sys
+import logging
+
+# Add src directory to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'templates', 'spaces', 'trackio'))
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+def test_dataset_manager():
+    """Test dataset manager functionality"""
+    try:
+        from dataset_utils import TrackioDatasetManager
+        
+        # Test with environment variables
+        hf_token = os.environ.get('HF_TOKEN')
+        dataset_repo = os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
+        
+        if not hf_token:
+            logger.warning("⚠️ HF_TOKEN not found in environment")
+            return False
+        
+        logger.info(f"🔧 Testing dataset manager with repo: {dataset_repo}")
+        
+        # Initialize dataset manager
+        manager = TrackioDatasetManager(dataset_repo, hf_token)
+        
+        # Test loading experiments
+        experiments = manager.load_existing_experiments()
+        logger.info(f"📊 Loaded {len(experiments)} experiments from dataset")
+        
+        # Test creating a sample experiment
+        sample_experiment = {
+            'experiment_id': f'test_diagnostic_{int(os.urandom(4).hex(), 16)}',
+            'name': 'Diagnostic Test Experiment',
+            'description': 'Test experiment created by diagnostic script',
+            'created_at': '2025-01-27T12:00:00',
+            'status': 'completed',
+            'metrics': '[]',
+            'parameters': '{"test": true}',
+            'artifacts': '[]',
+            'logs': '[]',
+            'last_updated': '2025-01-27T12:00:00'
+        }
+        
+        # Test upsert functionality
+        logger.info("🧪 Testing experiment upsert...")
+        success = manager.upsert_experiment(sample_experiment)
+        
+        if success:
+            logger.info("✅ Dataset manager working correctly")
+            
+            # Verify the experiment was saved
+            experiments_after = manager.load_existing_experiments()
+            logger.info(f"📊 After upsert: {len(experiments_after)} experiments")
+            
+            return True
+        else:
+            logger.error("❌ Failed to upsert test experiment")
+            return False
+            
+    except ImportError as e:
+        logger.error(f"❌ Failed to import dataset_utils: {e}")
+        return False
+    except Exception as e:
+        logger.error(f"❌ Dataset manager test failed: {e}")
+        return False
+
+def test_trackio_space():
+    """Test TrackioSpace initialization"""
+    try:
+        # Import the TrackioSpace class
+        from app import TrackioSpace
+        
+        logger.info("🧪 Testing TrackioSpace initialization...")
+        
+        # Initialize TrackioSpace
+        space = TrackioSpace()
+        
+        logger.info(f"📊 TrackioSpace initialized with {len(space.experiments)} experiments")
+        logger.info(f"🛡️ Dataset manager available: {'Yes' if space.dataset_manager else 'No'}")
+        logger.info(f"🔑 HF Token available: {'Yes' if space.hf_token else 'No'}")
+        logger.info(f"📂 Dataset repo: {space.dataset_repo}")
+        
+        return True
+        
+    except ImportError as e:
+        logger.error(f"❌ Failed to import TrackioSpace: {e}")
+        return False
+    except Exception as e:
+        logger.error(f"❌ TrackioSpace test failed: {e}")
+        return False
+
+def test_environment():
+    """Test environment configuration"""
+    logger.info("🔍 Checking environment configuration...")
+    
+    # Check required environment variables
+    env_vars = {
+        'HF_TOKEN': os.environ.get('HF_TOKEN'),
+        'TRACKIO_DATASET_REPO': os.environ.get('TRACKIO_DATASET_REPO'),
+        'TRACKIO_URL': os.environ.get('TRACKIO_URL'),
+        'SPACE_ID': os.environ.get('SPACE_ID')
+    }
+    
+    for var, value in env_vars.items():
+        if value:
+            masked_value = value[:8] + '...' if len(value) > 8 and 'TOKEN' in var else value
+            logger.info(f"✅ {var}: {masked_value}")
+        else:
+            logger.warning(f"⚠️ {var}: Not set")
+    
+    # Check if running on HF Spaces
+    is_hf_spaces = bool(os.environ.get('SPACE_ID'))
+    logger.info(f"🚀 Running on HF Spaces: {'Yes' if is_hf_spaces else 'No'}")
+    
+    return True
+
+def fix_common_issues():
+    """Suggest fixes for common issues"""
+    logger.info("💡 Common issue fixes:")
+    
+    # Check dataset repository format
+    dataset_repo = os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
+    if '/' not in dataset_repo:
+        logger.warning(f"⚠️ Dataset repo format issue: {dataset_repo} should be 'username/dataset-name'")
+    else:
+        logger.info(f"✅ Dataset repo format looks good: {dataset_repo}")
+    
+    # Check for URL issues
+    trackio_url = os.environ.get('TRACKIO_URL', 'https://tonic-test-trackio-test.hf.space')
+    if trackio_url.startswith('https://https://') or trackio_url.startswith('http://http://'):
+        logger.warning(f"⚠️ URL format issue detected: {trackio_url}")
+        fixed_url = trackio_url.replace('https://https://', 'https://').replace('http://http://', 'http://')
+        logger.info(f"💡 Fixed URL should be: {fixed_url}")
+    else:
+        logger.info(f"✅ Trackio URL format looks good: {trackio_url}")
+
+def main():
+    """Run all diagnostic tests"""
+    logger.info("🔧 Starting Trackio Space diagnostics...")
+    logger.info("=" * 60)
+    
+    try:
+        # Test environment
+        test_environment()
+        logger.info("-" * 40)
+        
+        # Test dataset manager
+        dataset_manager_ok = test_dataset_manager()
+        logger.info("-" * 40)
+        
+        # Test TrackioSpace
+        trackio_space_ok = test_trackio_space()
+        logger.info("-" * 40)
+        
+        # Suggest fixes
+        fix_common_issues()
+        logger.info("-" * 40)
+        
+        # Summary
+        logger.info("📋 DIAGNOSTIC SUMMARY:")
+        logger.info(f"Dataset Manager: {'✅ OK' if dataset_manager_ok else '❌ Issues'}")
+        logger.info(f"TrackioSpace: {'✅ OK' if trackio_space_ok else '❌ Issues'}")
+        
+        if dataset_manager_ok and trackio_space_ok:
+            logger.info("🎉 All systems appear to be working correctly!")
+            logger.info("💡 The issues in the logs might be related to:")
+            logger.info("   - Empty dataset (expected for new setup)")
+            logger.info("   - API client URL formatting (being auto-fixed)")
+            logger.info("   - Remote data access (falling back to local data)")
+        else:
+            logger.warning("⚠️ Some issues detected. Check the logs above for details.")
+        
+    except Exception as e:
+        logger.error(f"❌ Diagnostic script failed: {e}")
+        return False
+    
+    return True
+
+if __name__ == "__main__":
+    main()