Upload 4 files

requirements.txt

@@ -1,14 +1,8 @@
 torch>=1.9.0
 numpy>=1.19.2
 pandas>=1.2.4
-nilearn>=0.8.1
-nibabel>=3.2.1
 scikit-learn>=0.24.2
 matplotlib>=3.4.2
-gradio>=2.0.0
-datasets>=1.11.0
-huggingface_hub>=0.15.0
-transformers>=4.15.0
-seaborn>=0.11.2
+gradio>=3.0.0
 joblib>=1.0.1
 

test_huggingface.py

@@ -0,0 +1,35 @@
+"""
+Simple test script to verify the Huggingface app works locally.
+This will run the app with synthetic data.
+"""
+import numpy as np
+import pandas as pd
+import os
+
+# Ensure directories exist
+os.makedirs('results', exist_ok=True)
+os.makedirs('models', exist_ok=True)
+
+# Create synthetic data
+print("Creating synthetic test data...")
+n_samples = 10
+n_features = 100
+
+# Create FC matrix data
+fc_data = np.random.randn(n_samples, n_features)
+np.save('results/test_fc.npy', fc_data)
+print(f"Saved FC matrix data to results/test_fc.npy with shape {fc_data.shape}")
+
+# Create demographics data
+demo_df = pd.DataFrame({
+    'age': np.random.normal(60, 10, n_samples),
+    'sex': np.random.choice(['M', 'F'], n_samples),
+    'months_post_stroke': np.random.normal(24, 12, n_samples),
+    'wab_score': np.random.normal(65, 15, n_samples)
+})
+demo_df.to_csv('results/test_demographics.csv', index=False)
+print(f"Saved demographics data to results/test_demographics.csv with shape {demo_df.shape}")
+
+print("\nTest data created successfully!")
+print("\nNow you can run: python app.py")
+print("Then upload the test files to train a model.")

vae_model.py

@@ -1,495 +1,355 @@
+"""
+Simplified VAE implementation with explicit loss tracking.
+"""
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-from utils import to_torch, to_cuda, to_numpy, demo_to_torch
+import os
+import matplotlib.pyplot as plt
 from sklearn.base import BaseEstimator
 
-class VAE(nn.Module):
-    def __init__(self, input_dim, latent_dim, demo_dim, use_cuda=True):
-        super(VAE, self).__init__()
+class SimpleVAE(nn.Module):
+    def __init__(self, input_dim, latent_dim, demo_dim):
+        super(SimpleVAE, self).__init__()
+        # Store dimensions
         self.input_dim = input_dim
         self.latent_dim = latent_dim
         self.demo_dim = demo_dim
-        self.use_cuda = use_cuda
         
-        # Create layers with standard parameters (no .float() call)
-        self.enc1 = nn.Linear(input_dim, 1000)
-        self.enc2 = nn.Linear(1000, latent_dim)
+        # Encoder (FC data → latent)
+        self.enc1 = nn.Linear(input_dim, 256)
+        self.enc2 = nn.Linear(256, latent_dim)
         
-        # Decoder
-        self.dec1 = nn.Linear(latent_dim+demo_dim, 1000)
-        self.dec2 = nn.Linear(1000, input_dim)
+        # Decoder (latent + demographics → FC reconstruction)
+        self.dec1 = nn.Linear(latent_dim + demo_dim, 256)
+        self.dec2 = nn.Linear(256, input_dim)
         
-        # Batch normalization layers
-        self.bn1 = nn.BatchNorm1d(1000)
-        self.bn2 = nn.BatchNorm1d(1000)
+    def encode(self, x):
+        """Encode FC data to latent space"""
+        h = F.relu(self.enc1(x))
+        return self.enc2(h)
+        
+    def decode(self, z, demo):
+        """Decode from latent space to FC reconstruction"""
+        # Combine latent with demographics
+        z_combined = torch.cat([z, demo], dim=1)
+        h = F.relu(self.dec1(z_combined))
+        return self.dec2(h)
         
-        # Move to CUDA if requested and available
-        if use_cuda and torch.cuda.is_available():
-            self.cuda()
+    def forward(self, x, demo):
+        """Full forward pass"""
+        z = self.encode(x)
+        return self.decode(z, demo)
 
-    def enc(self, x):
-        # First layer with activation
-        h = self.enc1(x)
-        h = F.relu(h)
+class DemoVAE:
+    def __init__(self, nepochs=50, batch_size=8, latent_dim=16, lr=1e-3):
+        """Simple VAE model with demographic conditioning"""
+        self.nepochs = nepochs
+        self.batch_size = batch_size
+        self.latent_dim = latent_dim
+        self.lr = lr
+        self.vae = None
+        self.train_losses = []
+        self.val_losses = []
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        # Apply batch norm - handle training vs eval mode automatically
-        h = self.bn1(h)
+    def preprocess_demo(self, demo_data, demo_types, n_samples=None):
+        """Process demographic data into one-hot encoded tensors"""
+        if n_samples is None:
+            n_samples = len(demo_data[0])
+            
+        processed_demos = []
+        total_dims = 0
         
-        # Output layer
-        z = self.enc2(h)
-        return z
-
-    def gen(self, n):
-        return to_cuda(torch.randn(n, self.latent_dim).float(), self.use_cuda)
-
-    def dec(self, z, demo):
-        # Concatenate latent code with demographic data
-        z_combined = to_cuda(torch.cat([z, demo], dim=1), self.use_cuda)
+        # Process each demographic variable
+        for i, (data, dtype) in enumerate(zip(demo_data, demo_types)):
+            if dtype == 'continuous':
+                # For continuous variables, just normalize
+                data_np = np.array(data).reshape(-1, 1)
+                mean, std = np.mean(data_np), np.std(data_np)
+                if std == 0:  # Handle constant values
+                    normalized = np.zeros_like(data_np)
+                else:
+                    normalized = (data_np - mean) / std
+                processed_demos.append(normalized)
+                total_dims += 1
+            elif dtype == 'categorical':
+                # For categorical, create one-hot encoding
+                data_list = list(data)
+                categories = sorted(list(set(data_list)))
+                
+                # Create one-hot vectors
+                one_hot = np.zeros((len(data_list), len(categories)))
+                for j, val in enumerate(data_list):
+                    idx = categories.index(val)
+                    one_hot[j, idx] = 1
+                    
+                processed_demos.append(one_hot)
+                total_dims += len(categories)
         
-        # First decoder layer with activation
-        h = self.dec1(z_combined)
-        h = F.relu(h)
+        # Combine all demographics
+        demo_tensor = np.hstack(processed_demos)
+        return torch.tensor(demo_tensor, dtype=torch.float32), total_dims
+    
+    def fit(self, X, demo_data, demo_types):
+        """Train the VAE model"""
+        # Convert to numpy arrays if needed
+        X = np.array(X)
         
-        # Apply batch norm - handle training vs eval mode automatically
-        h = self.bn2(h)
+        # Process demographics
+        print("Processing demographics...")
+        demo_tensor, demo_dim = self.preprocess_demo(demo_data, demo_types)
         
-        # Output layer
-        x = self.dec2(h)
-        return x
-
-class DemoVAE(BaseEstimator):
-    def __init__(self, **params):
-        self.set_params(**params)
-
-    @staticmethod
-    def get_default_params():
-        return dict(
-            latent_dim=32,
-            use_cuda=True,
-            nepochs=100,      # Changed from 1000 to 100 for faster testing 
-            pperiod=10,       # Changed from 100 to 10 to see more progress updates
-            bsize=5,          # Changed from 16 to 5 for small sample sizes
-            loss_C_mult=1,
-            loss_mu_mult=1,
-            loss_rec_mult=100,
-            loss_decor_mult=10,
-            loss_pred_mult=0.001,
-            alpha=100,
-            LR_C=100,
-            lr=1e-4,
-            weight_decay=0
-        )
-
-    def get_params(self, deep=True):
-        return {k: getattr(self, k) for k in self.get_default_params().keys()}
-
-    def set_params(self, **params):
-        for k, v in self.get_default_params().items():
-            setattr(self, k, params.get(k, v))
-        return self
-
-    def fit(self, x, demo, demo_types):
-        from utils import train_vae
-        
-        # Calculate demo_dim
-        demo_dim = 0
-        for d, t in zip(demo, demo_types):
-            if t == 'continuous':
-                demo_dim += 1
-            elif t == 'categorical':
-                demo_dim += len(set(d))
-            else:
-                raise ValueError(f'Demographic type "{t}" not supported')
-        
-        # Initialize VAE
-        self.input_dim = x.shape[1]
-        self.demo_dim = demo_dim
-        self.vae = VAE(self.input_dim, self.latent_dim, demo_dim, self.use_cuda)
-        
-        # Train VAE
-        train_losses, val_losses = train_vae(
-            self.vae, x, demo, demo_types,
-            self.nepochs, self.pperiod, self.bsize,
-            self.loss_C_mult, self.loss_mu_mult, self.loss_rec_mult,
-            self.loss_decor_mult, self.loss_pred_mult,
-            self.lr, self.weight_decay, self.alpha, self.LR_C,
-            self
-        )
+        # Initialize model
+        input_dim = X.shape[1]
+        print(f"Creating model with input_dim={input_dim}, latent_dim={self.latent_dim}, demo_dim={demo_dim}")
+        self.vae = SimpleVAE(input_dim, self.latent_dim, demo_dim)
+        self.vae.to(self.device)
         
-        # Store the losses for later visualization
-        self.train_losses = train_losses
-        self.val_losses = val_losses
+        # Convert data to tensors
+        X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device)
+        demo_tensor = demo_tensor.to(self.device)
         
-        # Return the losses for immediate use
-        return train_losses, val_losses
-
-    def transform(self, x, demo, demo_types):
-        """
-        Transform data through the VAE model.
-        
-        Args:
-            x: Either an integer (to generate samples) or input data to encode/decode
-            demo: Demographic data
-            demo_types: Types of demographic variables
-            
-        Returns:
-            Transformed data (reconstructions or generations)
-        """
-        print(f"VAE transform called - Input type: {type(x)}")
-        if not isinstance(x, int):
-            print(f"Input data shape: {np.array(x).shape}")
-        print(f"Demo data: {[len(d) for d in demo]}, Types: {demo_types}")
-        
-        # Set model to evaluation mode to handle batch norm with batch size of 1
+        # Initialize optimizer
+        optimizer = torch.optim.Adam(self.vae.parameters(), lr=self.lr)
+        
+        # Training loop
+        n_samples = X.shape[0]
+        batch_size = min(self.batch_size, n_samples)
+        
+        # Clear any old losses
+        self.train_losses = []
+        self.val_losses = []
+        
+        # Initial validation loss
         self.vae.eval()
+        with torch.no_grad():
+            reconstructed = self.vae(X_tensor, demo_tensor)
+            init_val_loss = F.mse_loss(reconstructed, X_tensor).item()
+            self.val_losses.append(init_val_loss)
+            print(f"Initial validation loss: {init_val_loss:.4f}")
         
-        try:
-            # Use torch.no_grad to disable gradient calculation during inference
-            with torch.no_grad():
-                # Generate latent vectors or encode inputs
-                if isinstance(x, int):
-                    print(f"Generating {x} random latent vectors...")
-                    z = self.vae.gen(x)
-                    print(f"Generated latent vectors shape: {z.shape}")
-                else:
-                    print("Encoding input data to latent space...")
-                    x_tensor = to_cuda(to_torch(x), self.vae.use_cuda)
-                    print(f"Input tensor shape: {x_tensor.shape}")
-                    z = self.vae.enc(x_tensor)
-                    print(f"Encoded latent vectors shape: {z.shape}")
-                
-                # Convert demographics to tensors
-                print("Converting demographics to tensors...")
-                try:
-                    demo_t = demo_to_torch(demo, demo_types, self.pred_stats, self.vae.use_cuda)
-                    print(f"Demographic tensor shape: {demo_t.shape}")
-                except Exception as demo_err:
-                    print(f"Error in demographic conversion: {demo_err}")
-                    raise
-                
-                # Handle batch size of 1 for batch normalization
-                print(f"Decoding with batch size: {z.size(0)}")
-                if z.size(0) == 1:
-                    print("Using special handling for batch size=1...")
-                    # If batch size is 1, we need to be careful with batch norm
-                    # Clone and repeat the input to create a fake batch if needed
-                    if hasattr(self.vae, 'bn1') or hasattr(self.vae, 'bn2'):
-                        print("Batch normalization layers detected")
-                        try:
-                            # Try normal decoding first
-                            print("Attempting normal decoding...")
-                            y = self.vae.dec(z, demo_t)
-                            print("Normal decoding succeeded")
-                        except Exception as e:
-                            # If it fails, use a workaround for batch norm
-                            print(f"Normal decoding failed: {e}")
-                            print("Using batch norm workaround (repeating batch)...")
-                            # Create a batch by repeating the input
-                            z_batch = z.repeat(2, 1)
-                            demo_t_batch = demo_t.repeat(2, 1)
-                            # Get the output and use only the first element
-                            print(f"Created batch with shapes - z: {z_batch.shape}, demo: {demo_t_batch.shape}")
-                            y_batch = self.vae.dec(z_batch, demo_t_batch)
-                            print(f"Batch decoding succeeded, extracting first item from {y_batch.shape}")
-                            y = y_batch[0:1]
-                    else:
-                        # No batch norm, proceed normally
-                        print("No batch norm, proceeding normally...")
-                        y = self.vae.dec(z, demo_t)
-                else:
-                    # Normal batch size, proceed as usual
-                    print("Normal batch size, proceeding with standard decoding...")
-                    y = self.vae.dec(z, demo_t)
-                
-                print(f"Decoding complete, output tensor shape: {y.shape}")
+        # Main training loop
+        for epoch in range(self.nepochs):
+            epoch_losses = []
+            self.vae.train()
+            
+            # Process in batches
+            for i in range(0, n_samples, batch_size):
+                # Get batch
+                end = min(i + batch_size, n_samples)
+                x_batch = X_tensor[i:end]
+                demo_batch = demo_tensor[i:end]
                 
-                # Convert to numpy
-                result = to_numpy(y)
-                print(f"Final output shape: {result.shape}")
+                # Forward pass
+                optimizer.zero_grad()
+                reconstructed = self.vae(x_batch, demo_batch)
                 
-                # Check for NaN values in the result
-                if np.any(np.isnan(result)):
-                    print("WARNING: Result contains NaN values")
-                    result = np.nan_to_num(result)
-                    print("NaN values replaced with zeros")
+                # Calculate loss
+                loss = F.mse_loss(reconstructed, x_batch)
                 
-                return result
+                # Backward pass
+                loss.backward()
+                optimizer.step()
                 
-        except Exception as e:
-            import traceback
-            print(f"Error in VAE transform: {e}")
-            print(f"Traceback: {traceback.format_exc()}")
+                # Record loss
+                epoch_losses.append(loss.item())
             
-            # Create a fallback output with appropriate shape
-            if isinstance(x, int):
-                # Generate empty latent vectors with the right shape
-                n_features = self.input_dim
-                fallback = np.zeros((x, n_features))
-            else:
-                # Return empty array with same shape as input
-                fallback = np.zeros_like(np.array(x))
-                
-            print(f"Returning fallback output with shape: {fallback.shape}")
-            return fallback
-
-    def encode(self, x):
-        """Alias for get_latents method - to provide compatibility with some interfaces"""
-        return self.get_latents(x)
+            # End of epoch
+            avg_loss = np.mean(epoch_losses)
+            self.train_losses.append(avg_loss)
+            
+            # Validation
+            self.vae.eval()
+            with torch.no_grad():
+                reconstructed = self.vae(X_tensor, demo_tensor)
+                val_loss = F.mse_loss(reconstructed, X_tensor).item()
+                self.val_losses.append(val_loss)
+            
+            # Print progress every few epochs
+            if (epoch + 1) % 5 == 0 or epoch == 0:
+                print(f"Epoch {epoch+1}/{self.nepochs} - "
+                      f"Train loss: {avg_loss:.4f}, Val loss: {val_loss:.4f}")
+        
+        print(f"Training complete! Final loss: {self.train_losses[-1]:.4f}")
+        print(f"Loss history: {len(self.train_losses)} train, {len(self.val_losses)} validation")
+        
+        return self.train_losses, self.val_losses
+    
+    def transform(self, X, demo_data, demo_types):
+        """Generate reconstructions or synthetic samples"""
+        # Check if model is available
+        if self.vae is None:
+            raise ValueError("Model not trained or loaded yet")
         
-    def get_latents(self, x):
         # Set model to evaluation mode
         self.vae.eval()
         
-        # Use torch.no_grad for inference
-        with torch.no_grad():
-            try:
-                # Convert to torch tensor and move to CUDA if needed
-                x_tensor = to_cuda(to_torch(x), self.vae.use_cuda)
-                
-                # Get latent representation
-                z = self.vae.enc(x_tensor)
-            except Exception as e:
-                print(f"Error in encoder: {e}")
-                # Try workaround for batch norm if needed
-                if x.shape[0] == 1 and (hasattr(self.vae, 'bn1') or hasattr(self.vae, 'bn2')):
-                    print("Using batch normalization workaround for single sample")
-                    # Repeat the input to create a batch of size 2
-                    if len(x.shape) == 2:
-                        x_batch = np.repeat(x, 2, axis=0)
-                    else:
-                        x_batch = np.array([x[0], x[0]])
-                    
-                    # Process the batch
-                    x_tensor = to_cuda(to_torch(x_batch), self.vae.use_cuda)
-                    z_batch = self.vae.enc(x_tensor)
-                    
-                    # Extract just the first sample's latent representation
-                    z = z_batch[0:1]
+        # Check if we're generating or reconstructing
+        if isinstance(X, int):
+            # Generating n random samples
+            n_samples = X
+            
+            # Process demo data (repeat single values if needed)
+            demo_list = []
+            for d in demo_data:
+                if not isinstance(d, (list, np.ndarray)):
+                    # Single value, repeat for all samples
+                    demo_list.append([d] * n_samples)
                 else:
-                    # Re-raise if we can't handle it
-                    raise
-        
-        return to_numpy(z)
-
-    def save(self, path):
-        train_losses = getattr(self, 'train_losses', [])
-        val_losses = getattr(self, 'val_losses', [])
-        
-        # Make sure train_losses and val_losses are regular Python lists of float
-        if train_losses:
-            train_losses = [float(x) for x in train_losses]
-        else:
-            train_losses = []
+                    demo_list.append(d)
             
-        if val_losses:
-            val_losses = [float(x) for x in val_losses]
-        else:
-            val_losses = []
-        
-        # Save state dict separately (most compatible way)
-        torch.save(self.vae.state_dict(), f"{path}_state_dict")
-        print(f"Saved VAE model state to {path}_state_dict")
-        
-        # Save metadata as simple numpy arrays
-        import numpy as np
-        import json
-        np.savez(
-            f"{path}_metadata.npz",
-            train_losses=np.array(train_losses, dtype=np.float32),
-            val_losses=np.array(val_losses, dtype=np.float32),
-            input_dim=np.array([self.input_dim], dtype=np.int32),
-            demo_dim=np.array([self.demo_dim], dtype=np.int32)
-        )
-        
-        # Save parameters and pred_stats to JSON
-        params_json = {}
-        for k, v in self.get_params().items():
-            if isinstance(v, (int, float)):
-                params_json[k] = float(v)
-            elif isinstance(v, bool):
-                params_json[k] = v
-            else:
-                params_json[k] = str(v)
-                
-        # Convert pred_stats to JSON-serializable format
-        pred_stats_json = []
-        for stat in self.pred_stats:
-            if isinstance(stat, (list, tuple)):
-                pred_stats_json.append([float(v) if isinstance(v, (int, float)) else str(v) for v in stat])
-            else:
-                pred_stats_json.append(stat)
-                
-        with open(f"{path}_params.json", 'w') as f:
-            json.dump({
-                'params': params_json,
-                'pred_stats': pred_stats_json
-            }, f)
-        
-        # Also save with original method as a backup
-        try:
-            model_dict = {
-                'model_state_dict': self.vae.state_dict(),
-                'params': params_json,
-                'pred_stats': pred_stats_json,
-                'input_dim': int(self.input_dim),
-                'demo_dim': int(self.demo_dim),
-                'train_losses': train_losses,
-                'val_losses': val_losses
-            }
-            torch.save(model_dict, path)
-            print(f"Saved VAE model to {path}")
-        except Exception as e:
-            print(f"Error saving model with default settings: {e}")
-            print(f"Falling back to component files {path}_*")
-
-    def load(self, path):
-        # Simplified load function focusing on component-based loading first
-        try:
-            print(f"Attempting to load model from component files {path}_*")
-            import json
-            import numpy as np
-            import os
+            print(f"Generating {n_samples} samples with demo data: {demo_list}")
             
-            # Check if component files exist
-            state_dict_path = f"{path}_state_dict"
-            metadata_path = f"{path}_metadata.npz"
-            params_path = f"{path}_params.json"
+            # Process demographics
+            demo_tensor, demo_dim = self.preprocess_demo(demo_list, demo_types, n_samples)
             
-            if os.path.exists(state_dict_path) and os.path.exists(metadata_path) and os.path.exists(params_path):
-                # Load state dict from the most reliable source
-                print(f"Loading state dict from {state_dict_path}")
-                state_dict = torch.load(state_dict_path, map_location='cpu')
-                
-                # Load metadata
-                print(f"Loading metadata from {metadata_path}")
-                metadata = np.load(metadata_path, allow_pickle=True)
-                self.input_dim = int(metadata['input_dim'][0])
-                self.demo_dim = int(metadata['demo_dim'][0])
-                
-                # Load training histories if available
-                if 'train_losses' in metadata:
-                    self.train_losses = metadata['train_losses'].tolist()
+            # Generate random latent vectors
+            z = torch.randn(n_samples, self.latent_dim).to(self.device)
+            
+        else:
+            # Reconstructing existing data
+            X = np.array(X)
+            n_samples = X.shape[0]
+            
+            # Process demo data (repeat single values if needed)
+            demo_list = []
+            for d in demo_data:
+                if not isinstance(d, (list, np.ndarray)) or len(d) != n_samples:
+                    # Single value, repeat for all samples
+                    demo_list.append([d] * n_samples)
                 else:
-                    self.train_losses = []
-                    
-                if 'val_losses' in metadata:
-                    self.val_losses = metadata['val_losses'].tolist()
+                    demo_list.append(d)
+            
+            # Process demographics
+            demo_tensor, demo_dim = self.preprocess_demo(demo_list, demo_types)
+            
+            # Encode input data
+            X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device)
+            z = self.vae.encode(X_tensor)
+        
+        # Print shapes for debugging    
+        print(f"Latent shape: {z.shape}, Demo tensor shape: {demo_tensor.shape}")
+        
+        # Decode to get output
+        demo_tensor = demo_tensor.to(self.device)
+        with torch.no_grad():
+            # Make sure demo_tensor has the right dimensions
+            if demo_tensor.shape[1] != self.vae.demo_dim:
+                print(f"WARNING: Demo dimension mismatch. Expected {self.vae.demo_dim}, got {demo_tensor.shape[1]}")
+                # Use demographic dimension from the model
+                if demo_tensor.shape[1] > self.vae.demo_dim:
+                    # Trim extra dimensions
+                    demo_tensor = demo_tensor[:, :self.vae.demo_dim]
                 else:
-                    self.val_losses = []
+                    # Pad with zeros
+                    padding = torch.zeros(demo_tensor.shape[0], self.vae.demo_dim - demo_tensor.shape[1]).to(self.device)
+                    demo_tensor = torch.cat([demo_tensor, padding], dim=1)
+                print(f"Adjusted demo tensor shape: {demo_tensor.shape}")
                 
-                # Load parameters and pred_stats
-                print(f"Loading parameters from {params_path}")
-                with open(params_path, 'r') as f:
-                    json_data = json.load(f)
-                    self.set_params(**json_data['params'])
-                    self.pred_stats = json_data['pred_stats']
-                
-                # Initialize model and load state dict
-                print("Initializing VAE model with loaded parameters")
-                try:
-                    # First create model with proper typing
-                    device = torch.device("cpu")  # Always start with CPU
-                    self.vae = VAE(
-                        input_dim=int(self.input_dim), 
-                        latent_dim=int(self.latent_dim), 
-                        demo_dim=int(self.demo_dim), 
-                        use_cuda=False  # Initially False, move to CUDA later if needed
-                    )
-                    
-                    # Then load state dict
-                    self.vae.load_state_dict(state_dict)
-                    print(f"Successfully created VAE model and loaded state dict")
-                    
-                    # Move to CUDA if needed
-                    if self.use_cuda and torch.cuda.is_available():
-                        self.vae.cuda()
-                        print("Moved model to CUDA")
-                except Exception as e:
-                    print(f"Error initializing VAE model: {e}")
-                    # Create model without trying to use saved parameters
-                    self.vae = VAE(
-                        input_dim=100,  # Default size
-                        latent_dim=16,  # Small default 
-                        demo_dim=4,     # Default
-                        use_cuda=False  # Avoid CUDA issues
-                    )
-                    print("Created default VAE model (loading state dict failed)")
-                
-                print(f"Successfully loaded VAE model from component files {path}_*")
-                
-            # If component files don't exist, try loading the combined file
-            else:
-                print(f"Component files not found. Trying to load from {path}")
-                try:
-                    # Simple approach for PyTorch 2.1
-                    checkpoint = torch.load(path, map_location='cpu')
-                    
-                    # Initialize from checkpoint
-                    self.set_params(**checkpoint['params'])
-                    self.pred_stats = checkpoint['pred_stats']
-                    self.input_dim = checkpoint['input_dim']
-                    self.demo_dim = checkpoint['demo_dim']
-                    
-                    # Initialize model and load state dict
-                    try:
-                        # Create model on CPU first
-                        self.vae = VAE(
-                            input_dim=int(self.input_dim), 
-                            latent_dim=int(self.latent_dim), 
-                            demo_dim=int(self.demo_dim), 
-                            use_cuda=False  # Start with CPU
-                        )
-                        
-                        # Then load state dict
-                        self.vae.load_state_dict(checkpoint['model_state_dict'])
-                        
-                        # Move to CUDA if needed
-                        if self.use_cuda and torch.cuda.is_available():
-                            self.vae.cuda()
-                    except Exception as e:
-                        print(f"Error creating VAE model: {e}")
-                        # Fallback to a default model
-                        self.vae = VAE(
-                            input_dim=100,
-                            latent_dim=16,
-                            demo_dim=4,
-                            use_cuda=False
-                        )
-                    
-                    # Load training history
-                    if 'train_losses' in checkpoint:
-                        self.train_losses = checkpoint['train_losses']
-                    if 'val_losses' in checkpoint:
-                        self.val_losses = checkpoint['val_losses']
-                        
-                    print(f"Successfully loaded VAE model from {path}")
-                except Exception as e:
-                    print(f"Error loading model: {e}")
-                    raise
-        except Exception as e:
-            import os
-            print(f"Error during model loading: {e}")
-            print("Available files in models directory:")
-            if os.path.exists('models'):
-                print('\n'.join(os.listdir('models')))
-            else:
-                print("models directory does not exist")
+            output = self.vae.decode(z, demo_tensor)
             
-            # Create a minimal model for fallback
-            print("Creating a new untrained model as fallback")
-            self.input_dim = 100  # Default size for a typical FC matrix
-            self.demo_dim = 4     # Default for common demographic variables
-            self.pred_stats = []
-            self.train_losses = []
-            self.val_losses = []
-            self.vae = VAE(self.input_dim, self.latent_dim, self.demo_dim, self.use_cuda)
+        # Convert to numpy
+        return output.cpu().numpy()
+    
+    def get_latents(self, X):
+        """Encode data to latent representations"""
+        X = np.array(X)
+        X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device)
+        
+        with torch.no_grad():
+            z = self.vae.encode(X_tensor)
+        
+        return z.cpu().numpy()
+        
+    def save(self, path):
+        """Save the model and training history"""
+        # Ensure the directory exists
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+        
+        # Create state dict with all necessary info
+        state = {
+            'vae_state': self.vae.state_dict(),
+            'input_dim': self.vae.input_dim,
+            'latent_dim': self.latent_dim,
+            'demo_dim': self.vae.demo_dim,
+            'train_losses': self.train_losses,
+            'val_losses': self.val_losses,
+            'nepochs': self.nepochs,
+            'batch_size': self.batch_size,
+            'lr': self.lr
+        }
+        
+        # Save the model
+        torch.save(state, path)
+        print(f"Model saved to {path}")
+        
+        # Print info about saved losses
+        print(f"Saved loss data: {len(self.train_losses)} train, {len(self.val_losses)} validation")
+        
+    def load(self, path):
+        """Load the model from a file"""
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Model file not found: {path}")
             
-            raise RuntimeError(f"Unable to load VAE model: {e}")
+        # Load state dict
+        state = torch.load(path, map_location=self.device)
+        
+        # Set attributes
+        self.latent_dim = state['latent_dim']
+        self.nepochs = state.get('nepochs', 50)
+        self.batch_size = state.get('batch_size', 8)
+        self.lr = state.get('lr', 1e-3)
+        self.train_losses = state.get('train_losses', [])
+        self.val_losses = state.get('val_losses', [])
         
-        # Move model to appropriate device after loading
-        if self.use_cuda and torch.cuda.is_available():
-            self.vae.cuda()
+        # Create model
+        self.vae = SimpleVAE(
+            input_dim=state['input_dim'],
+            latent_dim=self.latent_dim,
+            demo_dim=state['demo_dim']
+        )
+        
+        # Load weights
+        self.vae.load_state_dict(state['vae_state'])
+        self.vae.to(self.device)
+        
+        print(f"Model loaded from {path}")
+        print(f"Loaded loss data: {len(self.train_losses)} train, {len(self.val_losses)} validation")
+        
+def plot_learning_curves(train_losses, val_losses):
+    """Plot training and validation loss curves"""
+    # Create figure
+    plt.figure(figsize=(10, 6))
+    
+    # Check if we have loss data
+    if not train_losses:
+        plt.text(0.5, 0.5, "No training loss data available", 
+                ha='center', va='center', transform=plt.gca().transAxes,
+                fontsize=14, color='red')
+        plt.axis('off')
+        return plt.gcf()
+    
+    # Plot losses
+    epochs = range(1, len(train_losses) + 1)
+    plt.plot(epochs, train_losses, 'b-', label='Training loss')
+    
+    if val_losses:
+        # Adjust validation epochs if lengths differ
+        if len(val_losses) == len(train_losses) + 1:
+            # Initial validation + epoch validations
+            val_epochs = [0] + list(epochs)
         else:
-            self.vae.cpu()
+            val_epochs = epochs[:len(val_losses)]
+            
+        plt.plot(val_epochs, val_losses, 'r-', label='Validation loss')
+    
+    # Add labels
+    plt.title('VAE Training and Validation Loss')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.legend()
+    plt.grid(True, alpha=0.3)
+    
+    return plt.gcf()