featrixsphere 0.2.3737__py3-none-any.whl → 0.2.4983__py3-none-any.whl

featrixsphere/__init__.py

@@ -38,7 +38,7 @@ Example:
     ...                                labels=['Experiment A', 'Experiment B'])
 """
 
-__version__ = "0.2.3737"
+__version__ = "0.2.4983"
 __author__ = "Featrix"
 __email__ = "support@featrix.com"
 __license__ = "MIT"

featrixsphere/client.py

@@ -11,7 +11,7 @@ import time
 import requests
 from pathlib import Path
 from typing import Dict, Any, Optional, List, Tuple, Union
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import gzip
 import os
 import random
@@ -65,6 +65,7 @@ class SessionInfo:
     status: str
     jobs: Dict[str, Any]
     job_queue_positions: Dict[str, Any]
+    job_plan: List[Dict[str, Any]] = field(default_factory=list)
     _client: Optional['FeatrixSphereClient'] = None
     
     def predictors(self) -> List[Dict[str, Any]]:
@@ -82,6 +83,51 @@ class SessionInfo:
             return list(predictors_dict.values())
         except Exception:
             return []
+    
+    def embedding_space_info(self) -> Optional[Dict[str, Any]]:
+        """
+        Get embedding space information for this session.
+        
+        Returns:
+            Dictionary with ES info (dimensions, epochs, etc.) or None if not available
+        """
+        if not self._client:
+            return None
+        
+        try:
+            # Get session details from the client
+            session_data = self._client._get_json(f"/compute/session/{self.session_id}")
+            
+            es_info = {}
+            
+            # Extract embedding space path
+            embedding_space_path = session_data.get('embedding_space')
+            if embedding_space_path:
+                es_info['embedding_space_path'] = embedding_space_path
+            
+            # Extract model architecture info
+            model_info = session_data.get('model_info', {}) or session_data.get('embedding_space', {})
+            if isinstance(model_info, dict):
+                es_info['d_model'] = model_info.get('d_model') or model_info.get('embedding_dim')
+                es_info['dimensions'] = es_info.get('d_model')  # Alias for compatibility
+                es_info['parameter_count'] = model_info.get('parameter_count') or model_info.get('num_parameters')
+                es_info['layer_count'] = model_info.get('layer_count') or model_info.get('num_layers')
+            
+            # Extract training statistics
+            training_stats = session_data.get('training_stats', {}) or session_data.get('stats', {})
+            if isinstance(training_stats, dict):
+                es_info['epochs'] = training_stats.get('final_epoch') or training_stats.get('epochs_trained') or training_stats.get('epochs')
+                es_info['final_loss'] = training_stats.get('final_loss') or training_stats.get('loss')
+                es_info['final_val_loss'] = training_stats.get('final_val_loss') or training_stats.get('validation_loss')
+                es_info['training_time_seconds'] = training_stats.get('training_time') or training_stats.get('elapsed_seconds')
+            
+            # If we have any info, return it
+            if es_info:
+                return es_info
+            
+            return None
+        except Exception:
+            return None
 
 
 class PredictionBatch:
@@ -617,6 +663,7 @@ class FeatrixSphereClient:
             status=response_data.get('status', 'unknown'),
             jobs={},
             job_queue_positions={},
+            job_plan=[],
             _client=self
         )
     
@@ -640,6 +687,7 @@ class FeatrixSphereClient:
         session = response_data.get('session', {})
         jobs = response_data.get('jobs', {})
         positions = response_data.get('job_queue_positions', {})
+        job_plan = session.get('job_plan', [])
         
         return SessionInfo(
             session_id=session.get('session_id', session_id),
@@ -647,6 +695,7 @@ class FeatrixSphereClient:
             status=session.get('status', 'unknown'),
             jobs=jobs,
             job_queue_positions=positions,
+            job_plan=job_plan,
             _client=self
         )
     
@@ -1370,6 +1419,20 @@ class FeatrixSphereClient:
         
         start_time = time.time()
         
+        # Initial wait for job dispatch (jobs are dispatched asynchronously after session creation)
+        # Wait up to 10 seconds for jobs to appear before starting main monitoring loop
+        initial_wait_timeout = 10
+        initial_wait_start = time.time()
+        jobs_appeared = False
+        
+        while time.time() - initial_wait_start < initial_wait_timeout:
+            session_info = self.get_session_status(session_id)
+            if session_info.jobs:
+                jobs_appeared = True
+                break
+            time.sleep(0.5)  # Check every 500ms during initial wait
+        
+        # Main monitoring loop
         while time.time() - start_time < max_wait_time:
             session_info = self.get_session_status(session_id)
             elapsed = time.time() - start_time
@@ -1842,7 +1905,8 @@ class FeatrixSphereClient:
             session_type=response_data.get('session_type', 'embedding_space'),
             status=response_data.get('status', 'ready'),
             jobs={},
-            job_queue_positions={}
+            job_queue_positions={},
+            job_plan=[]
         )
 
     def fine_tune_embedding_space(
@@ -1972,7 +2036,144 @@ class FeatrixSphereClient:
             session_type=response_data.get('session_type', 'embedding_space_finetune'),
             status=response_data.get('status', 'ready'),
             jobs={},
-            job_queue_positions={}
+            job_queue_positions={},
+            job_plan=[]
+        )
+
+    def extend_embedding_space(
+        self,
+        name: str,
+        parent_session_id: str = None,
+        parent_embedding_space_path: str = None,
+        s3_training_dataset: str = None,
+        s3_validation_dataset: str = None,
+        n_epochs: int = None,
+        webhooks: Dict[str, str] = None,
+        user_metadata: Dict[str, Any] = None
+    ) -> SessionInfo:
+        """
+        Extend an existing embedding space with new feature columns.
+        
+        This method takes a pre-trained embedding space and extends it with new feature columns
+        from enriched training/validation data. The extended ES preserves existing encoder weights
+        and creates new codecs for the new columns.
+        
+        **When to Use Extend vs Fine-Tune:**
+        - **Extend**: When you've added NEW COLUMNS (features) to your dataset
+        - **Fine-Tune**: When you have new rows with the SAME COLUMNS
+        
+        **How It Works:**
+        1. Loads the parent embedding space
+        2. Identifies new columns in the enriched dataset
+        3. Creates codecs for the new columns
+        4. Copies existing encoder weights (preserves learned representations)
+        5. Fine-tunes for shorter duration (default: original_epochs / 4)
+        6. Returns extended embedding space with all columns
+        
+        Args:
+            name: Name for the extended embedding space
+            parent_session_id: Session ID of the parent embedding space (optional)
+            parent_embedding_space_path: Direct path to parent embedding space pickle file (optional)
+            s3_training_dataset: S3 URL for enriched training dataset with new columns (must start with 's3://')
+            s3_validation_dataset: S3 URL for enriched validation dataset with new columns (must start with 's3://')
+            n_epochs: Number of epochs for extension training (default: original_epochs / 4)
+            webhooks: Optional dict with webhook configuration keys
+            user_metadata: Optional metadata dict to attach to the session
+            
+        Returns:
+            SessionInfo for the newly created extension session
+            
+        Raises:
+            ValueError: If S3 URLs are invalid or neither parent identifier is provided
+            
+        Example:
+            ```python
+            # Extend an existing embedding space with new feature columns
+            client = FeatrixSphereClient("https://sphere-api.featrix.com")
+            
+            # Original ES was trained on: age, income, credit_score
+            # New data includes engineered features: debt_to_income_ratio, age_bin
+            extended = client.extend_embedding_space(
+                name="customer_model_with_features",
+                parent_session_id="abc123-20240101-120000",
+                s3_training_dataset="s3://my-bucket/enriched_training.csv",
+                s3_validation_dataset="s3://my-bucket/enriched_validation.csv",
+                n_epochs=25  # Optional: specify epochs (defaults to original/4)
+            )
+            
+            # Wait for extension to complete
+            client.wait_for_session_completion(extended.session_id)
+            
+            # The extended ES now includes the new feature columns
+            # Use it for predictions with enriched data
+            result = client.predict(extended.session_id, {
+                "age": 35,
+                "income": 75000,
+                "credit_score": 720,
+                "debt_to_income_ratio": 0.25,  # New feature!
+                "age_bin": "30-40"  # New feature!
+            })
+            ```
+        """
+        # Validate S3 URLs
+        if s3_training_dataset and not s3_training_dataset.startswith('s3://'):
+            raise ValueError("s3_training_dataset must be a valid S3 URL (s3://...)")
+        if s3_validation_dataset and not s3_validation_dataset.startswith('s3://'):
+            raise ValueError("s3_validation_dataset must be a valid S3 URL (s3://...)")
+        
+        # Validate that we have either parent_session_id or parent_embedding_space_path
+        if not parent_session_id and not parent_embedding_space_path:
+            raise ValueError("Either parent_session_id or parent_embedding_space_path must be provided")
+        
+        print(f"Extending embedding space '{name}' with new features...")
+        if parent_session_id:
+            print(f"  Parent session: {parent_session_id}")
+        if parent_embedding_space_path:
+            print(f"  Parent embedding space: {parent_embedding_space_path}")
+        print(f"  Enriched training data: {s3_training_dataset}")
+        print(f"  Enriched validation data: {s3_validation_dataset}")
+        if n_epochs:
+            print(f"  Extension epochs: {n_epochs}")
+        else:
+            print(f"  Extension epochs: auto (original/4)")
+        
+        data = {
+            "name": name,
+            "s3_file_data_set_training": s3_training_dataset,
+            "s3_file_data_set_validation": s3_validation_dataset
+        }
+        
+        if parent_session_id:
+            data["parent_session_id"] = parent_session_id
+        if parent_embedding_space_path:
+            data["parent_embedding_space_path"] = parent_embedding_space_path
+        if n_epochs is not None:
+            data["n_epochs"] = n_epochs
+        
+        if webhooks:
+            data['webhooks'] = webhooks
+        if user_metadata:
+            import json
+            data['user_metadata'] = json.dumps(user_metadata)
+            print(f"User metadata: {user_metadata}")
+        
+        response_data = self._post_json("/compute/extend-embedding-space", data)
+        
+        session_id = response_data.get('session_id')
+        extend_info = response_data.get('extend_es_info', {})
+        
+        print(f"Extension session created: {session_id}")
+        if extend_info:
+            print(f"  Original epochs: {extend_info.get('original_epochs', 'N/A')}")
+            print(f"  Extension epochs: {extend_info.get('extension_epochs', 'N/A')}")
+        
+        return SessionInfo(
+            session_id=session_id,
+            session_type=response_data.get('session_type', 'embedding_space_extend'),
+            status=response_data.get('status', 'ready'),
+            jobs={},
+            job_queue_positions={},
+            job_plan=[]
         )
 
     # =========================================================================
@@ -2033,7 +2234,8 @@ class FeatrixSphereClient:
             session_type=response_data.get('session_type', 'sphere'),
             status=response_data.get('status', 'ready'),
             jobs={},
-            job_queue_positions={}
+            job_queue_positions={},
+            job_plan=[]
         )
 
     def upload_df_and_create_session(self, df=None, filename: str = "data.csv", file_path: str = None, 
@@ -2309,23 +2511,57 @@ class FeatrixSphereClient:
                 compression_ratio = (1 - compressed_size / original_size) * 100
                 print(f"Converted Parquet to CSV and compressed from {original_size:,} to {compressed_size:,} bytes ({compression_ratio:.1f}% reduction)")
             else:
-                # Regular CSV file - read and compress it
+                # Regular CSV file - check size and suggest Parquet for large files
                 with open(file_path, 'rb') as f:
                     csv_content = f.read()
                 
-                # Compress the content
-                print("Compressing CSV file...")
-                compressed_buffer = io.BytesIO()
-                with gzip.GzipFile(fileobj=compressed_buffer, mode='wb') as gz:
-                    gz.write(csv_content)
-                file_content = compressed_buffer.getvalue()
-                upload_filename = os.path.basename(file_path) + '.gz'
-                content_type = 'application/gzip'
+                csv_size_mb = len(csv_content) / (1024 * 1024)
+                CSV_WARNING_THRESHOLD_MB = 1.0  # Warn if CSV > 1MB
                 
-                original_size = len(csv_content)
-                compressed_size = len(file_content)
-                compression_ratio = (1 - compressed_size / original_size) * 100
-                print(f"Compressed from {original_size:,} to {compressed_size:,} bytes ({compression_ratio:.1f}% reduction)")
+                if csv_size_mb > CSV_WARNING_THRESHOLD_MB:
+                    print(f"\n⚠️  Warning: CSV file is {csv_size_mb:.1f} MB")
+                    print(f"   Parquet format is more efficient for large files (smaller size, faster upload).")
+                    print(f"   Converting to Parquet format for better performance...")
+                    
+                    # Read CSV as DataFrame
+                    csv_df = pd.read_csv(file_path)
+                    
+                    # Convert to Parquet in memory
+                    parquet_buffer = io.BytesIO()
+                    try:
+                        # Try pyarrow first (faster), fallback to fastparquet
+                        csv_df.to_parquet(parquet_buffer, index=False, engine='pyarrow')
+                    except (ImportError, ValueError):
+                        # Fallback to fastparquet or default engine
+                        try:
+                            csv_df.to_parquet(parquet_buffer, index=False, engine='fastparquet')
+                        except (ImportError, ValueError):
+                            # Last resort: use default engine
+                            csv_df.to_parquet(parquet_buffer, index=False)
+                    parquet_content = parquet_buffer.getvalue()
+                    parquet_size_mb = len(parquet_content) / (1024 * 1024)
+                    
+                    # Use Parquet instead of compressed CSV
+                    file_content = parquet_content
+                    upload_filename = os.path.basename(file_path).replace('.csv', '.parquet')
+                    content_type = 'application/octet-stream'
+                    
+                    size_reduction = (1 - len(parquet_content) / len(csv_content)) * 100
+                    print(f"   ✅ Converted to Parquet: {csv_size_mb:.1f} MB → {parquet_size_mb:.1f} MB ({size_reduction:.1f}% reduction)")
+                else:
+                    # Small CSV - compress as before
+                    print("Compressing CSV file...")
+                    compressed_buffer = io.BytesIO()
+                    with gzip.GzipFile(fileobj=compressed_buffer, mode='wb') as gz:
+                        gz.write(csv_content)
+                    file_content = compressed_buffer.getvalue()
+                    upload_filename = os.path.basename(file_path) + '.gz'
+                    content_type = 'application/gzip'
+                    
+                    original_size = len(csv_content)
+                    compressed_size = len(file_content)
+                    compression_ratio = (1 - compressed_size / original_size) * 100
+                    print(f"Compressed from {original_size:,} to {compressed_size:,} bytes ({compression_ratio:.1f}% reduction)")
         
         # Handle DataFrame input
         else:
@@ -2334,29 +2570,31 @@ class FeatrixSphereClient:
             
             print(f"Uploading DataFrame ({len(df)} rows, {len(df.columns)} columns)")
             
-            # Clean NaN values in DataFrame before CSV conversion
+            # Clean NaN values in DataFrame before conversion
             # This prevents JSON encoding issues when the server processes the data
             # Use pandas.notna() with where() for compatibility with all pandas versions
             cleaned_df = df.where(pd.notna(df), None)  # Replace NaN with None for JSON compatibility
             
-            # Convert DataFrame to CSV and compress
-            csv_buffer = io.StringIO()
-            cleaned_df.to_csv(csv_buffer, index=False)
-            csv_data = csv_buffer.getvalue().encode('utf-8')
-            
-            # Compress the CSV data
-            print("Compressing DataFrame...")
-            compressed_buffer = io.BytesIO()
-            with gzip.GzipFile(fileobj=compressed_buffer, mode='wb') as gz:
-                gz.write(csv_data)
-            file_content = compressed_buffer.getvalue()
-            upload_filename = filename if filename.endswith('.gz') else filename + '.gz'
-            content_type = 'application/gzip'
-            
-            original_size = len(csv_data)
-            compressed_size = len(file_content)
-            compression_ratio = (1 - compressed_size / original_size) * 100
-            print(f"Compressed from {original_size:,} to {compressed_size:,} bytes ({compression_ratio:.1f}% reduction)")
+            # Always use Parquet format for DataFrames (smaller, faster than CSV.gz)
+            print("Converting DataFrame to Parquet format...")
+            parquet_buffer = io.BytesIO()
+            try:
+                # Try pyarrow first (faster), fallback to fastparquet
+                cleaned_df.to_parquet(parquet_buffer, index=False, engine='pyarrow')
+            except (ImportError, ValueError):
+                # Fallback to fastparquet or default engine
+                try:
+                    cleaned_df.to_parquet(parquet_buffer, index=False, engine='fastparquet')
+                except (ImportError, ValueError):
+                    # Last resort: use default engine
+                    cleaned_df.to_parquet(parquet_buffer, index=False)
+            
+            file_content = parquet_buffer.getvalue()
+            parquet_size_mb = len(file_content) / (1024 * 1024)
+            upload_filename = filename.replace('.csv', '.parquet') if filename.endswith('.csv') else filename + '.parquet'
+            content_type = 'application/octet-stream'
+            
+            print(f"✅ Saved as Parquet: {parquet_size_mb:.2f} MB")
         
         # Upload the compressed file with optional column overrides
         files = {'file': (upload_filename, file_content, content_type)}
@@ -2403,9 +2641,18 @@ class FeatrixSphereClient:
         file_size_mb = len(file_content) / (1024 * 1024)
         CHUNK_SIZE_MB = 512  # 512 MB chunk size
         CHUNK_SIZE_BYTES = CHUNK_SIZE_MB * 1024 * 1024
+        LARGE_FILE_WARNING_MB = 10  # Warn if file > 10 MB
+        
+        if file_size_mb > LARGE_FILE_WARNING_MB:
+            print(f"\n⚠️  Warning: File size ({file_size_mb:.1f} MB) is quite large")
+            print(f"   For very large files (>10 MB), consider using S3 uploads:")
+            print(f"   1. Upload your file to S3 (or your cloud storage)")
+            print(f"   2. Generate a signed/private URL with read access")
+            print(f"   3. Contact Featrix support to configure S3-based uploads")
+            print(f"   This can be more reliable than direct uploads for large datasets.")
         
         if file_size_mb > CHUNK_SIZE_MB:
-            print(f"⚠️  Warning: File size ({file_size_mb:.1f} MB) exceeds {CHUNK_SIZE_MB} MB threshold")
+            print(f"\n⚠️  Warning: File size ({file_size_mb:.1f} MB) exceeds {CHUNK_SIZE_MB} MB threshold")
             print(f"   Large uploads may timeout. Consider splitting the data or using smaller batches.")
         
         # Try upload with retry on 504
@@ -2451,7 +2698,8 @@ class FeatrixSphereClient:
             session_type=response_data.get('session_type', 'sphere'),
             status=response_data.get('status', 'ready'),
             jobs={},
-            job_queue_positions={}
+            job_queue_positions={},
+            job_plan=[]
         )
         
 
@@ -2564,7 +2812,8 @@ class FeatrixSphereClient:
     # =========================================================================
     
     def predict(self, session_id: str, record: Dict[str, Any], target_column: str = None, 
-               predictor_id: str = None, max_retries: int = None, queue_batches: bool = False) -> Dict[str, Any]:
+               predictor_id: str = None, best_metric_preference: str = None,
+               max_retries: int = None, queue_batches: bool = False) -> Dict[str, Any]:
         """
         Make a single prediction for a record.
         
@@ -2573,6 +2822,7 @@ class FeatrixSphereClient:
             record: Record dictionary (without target column)
             target_column: Specific target column predictor to use (required if multiple predictors exist and predictor_id not specified)
             predictor_id: Specific predictor ID to use (recommended - more precise than target_column)
+            best_metric_preference: Which metric checkpoint to use: "roc_auc", "pr_auc", or None (use default checkpoint) (default: None)
             max_retries: Number of retries for errors (default: uses client default)
             queue_batches: If True, queue this prediction for batch processing instead of immediate API call
             
@@ -2595,41 +2845,24 @@ class FeatrixSphereClient:
             queue_id = self._add_to_prediction_queue(session_id, record, target_column, predictor_id)
             return {"queued": True, "queue_id": queue_id}
         
-        # Check if multiple predictors exist and require specification
-        predictors = self._get_available_predictors(session_id, debug=False)
-        if len(predictors) > 1 and not target_column and not predictor_id:
-            available_targets = [p.get('target_column') for p in predictors.values() if p.get('target_column')]
-            available_ids = list(predictors.keys())
-            raise ValueError(
-                f"Session {session_id} has {len(predictors)} predictors. "
-                f"You must specify either 'target_column' or 'predictor_id'.\n"
-                f"Available target columns: {available_targets}\n"
-                f"Available predictor IDs: {available_ids}\n"
-                f"Use client.list_predictors('{session_id}') to see details."
-            )
-        
-        # Resolve predictor information (handles both predictor_id and target_column)
-        predictor_info = self._resolve_predictor_id(session_id, predictor_id, target_column)
-        validated_target_column = predictor_info['target_column']
-        resolved_predictor_id = predictor_info['predictor_id']
-        
-        # Clean NaN/Inf values and remove target column
+        # Clean NaN/Inf values 
         cleaned_record = self._clean_numpy_values(record)
-        # Additional NaN cleaning for JSON encoding
         cleaned_record = self.replace_nans_with_nulls(cleaned_record)
-        cleaned_records = self._remove_target_columns(session_id, [cleaned_record], validated_target_column)
-        final_record = cleaned_records[0] if cleaned_records else cleaned_record
         
-        # Add predictor info to request so server knows exactly which predictor to use
+        # Build request payload - let the server handle predictor resolution
         request_payload = {
-            "query_record": final_record,
-            "target_column": validated_target_column
+            "query_record": cleaned_record,
         }
         
-        # Include predictor_id if available for server-side routing
-        if resolved_predictor_id:
-            request_payload["predictor_id"] = resolved_predictor_id
+        # Include whatever the caller provided - server will figure it out
+        if target_column:
+            request_payload["target_column"] = target_column
+        if predictor_id:
+            request_payload["predictor_id"] = predictor_id
+        if best_metric_preference:
+            request_payload["best_metric_preference"] = best_metric_preference
         
+        # Just send it to the server - it has all the smart fallback logic
         response_data = self._post_json(f"/session/{session_id}/predict", request_payload, max_retries=max_retries)
         return response_data
     
@@ -2712,7 +2945,7 @@ class FeatrixSphereClient:
     def plot_training_loss(self, session_id: str, figsize: Tuple[int, int] = (12, 8), 
                           style: str = 'notebook', save_path: Optional[str] = None,
                           show_learning_rate: bool = True, smooth: bool = True,
-                          title: Optional[str] = None) -> 'plt.Figure':
+                          title: Optional[str] = None):
         """
         Plot comprehensive training loss curves for a session (both embedding space and single predictor).
         
@@ -2800,7 +3033,7 @@ class FeatrixSphereClient:
 
     def plot_embedding_space_training(self, session_id: str, figsize: Tuple[int, int] = (10, 6),
                                      style: str = 'notebook', save_path: Optional[str] = None,
-                                     show_mutual_info: bool = False) -> 'plt.Figure':
+                                     show_mutual_info: bool = False):
         """
         Plot detailed embedding space training metrics.
         
@@ -2876,7 +3109,7 @@ class FeatrixSphereClient:
 
     def plot_single_predictor_training(self, session_id: str, figsize: Tuple[int, int] = (10, 6),
                                       style: str = 'notebook', save_path: Optional[str] = None,
-                                      show_metrics: bool = True) -> 'plt.Figure':
+                                      show_metrics: bool = True):
         """
         Plot detailed single predictor training metrics.
         
@@ -2952,7 +3185,7 @@ class FeatrixSphereClient:
 
     def plot_training_comparison(self, session_ids: List[str], labels: Optional[List[str]] = None,
                                figsize: Tuple[int, int] = (12, 8), style: str = 'notebook',
-                               save_path: Optional[str] = None) -> plt.Figure:
+                               save_path: Optional[str] = None):
         """
         Compare training curves across multiple sessions.
         
@@ -3205,7 +3438,7 @@ class FeatrixSphereClient:
     def plot_embedding_space_3d(self, session_id: str, sample_size: int = 2000,
                                 color_by: Optional[str] = None, size_by: Optional[str] = None,
                                 interactive: bool = True, style: str = 'notebook',
-                                title: Optional[str] = None, save_path: Optional[str] = None) -> Union[plt.Figure, 'go.Figure']:
+                                title: Optional[str] = None, save_path: Optional[str] = None):
         """
         Create interactive 3D visualization of the embedding space.
         
@@ -3281,7 +3514,7 @@ class FeatrixSphereClient:
                            style: str = 'notebook', save_path: Optional[str] = None,
                            show_embedding_evolution: bool = True, 
                            show_loss_evolution: bool = True,
-                           fps: int = 2, notebook_mode: bool = True) -> Union[plt.Figure, 'HTML']:
+                           fps: int = 2, notebook_mode: bool = True):
         """
         Create an animated training movie showing loss curves and embedding evolution.
         
@@ -3335,7 +3568,7 @@ class FeatrixSphereClient:
 
     def plot_embedding_evolution(self, session_id: str, epoch_range: Optional[Tuple[int, int]] = None,
                                  interactive: bool = True, sample_size: int = 1000,
-                                 color_by: Optional[str] = None) -> Union[plt.Figure, 'go.Figure']:
+                                 color_by: Optional[str] = None):
         """
         Show how embedding space evolves during training across epochs.
         
@@ -3771,7 +4004,18 @@ class FeatrixSphereClient:
         available_predictors = self._get_available_predictors(session_id, debug=debug)
         
         if not available_predictors:
-            raise ValueError(f"No trained predictors found for session {session_id}")
+            # Don't fail here - let the server try to find/auto-discover the predictor
+            # The server's /predict endpoint has smart fallback logic to find checkpoint files
+            # even if the session file wasn't properly updated (e.g., training crashed)
+            if debug:
+                print(f"⚠️  No predictors found via models endpoint, letting server handle discovery")
+            return {
+                'target_column': target_column,
+                'predictor_id': predictor_id,
+                'path': None,
+                'type': None,
+                'server_discovery': True  # Flag that server should auto-discover
+            }
         
         # If predictor_id is provided, find it directly (since it's now the key)
         if predictor_id:
@@ -4649,6 +4893,7 @@ class FeatrixSphereClient:
                 status="running",
                 jobs={},
                 job_queue_positions={},
+                job_plan=[],
                 _client=self
             )
 
@@ -5500,6 +5745,7 @@ class FeatrixSphereClient:
                     status="running",
                     jobs={},
                     job_queue_positions={},
+                    job_plan=[],
                     _client=self
                 )
                 
@@ -6083,7 +6329,8 @@ class FeatrixSphereClient:
     # =========================================================================
     
     def predict_table(self, session_id: str, table_data: Dict[str, Any], 
-                     target_column: str = None, predictor_id: str = None, max_retries: int = None) -> Dict[str, Any]:
+                     target_column: str = None, predictor_id: str = None, 
+                     best_metric_preference: str = None, max_retries: int = None) -> Dict[str, Any]:
         """
         Make batch predictions using JSON Tables format.
         
@@ -6124,6 +6371,8 @@ class FeatrixSphereClient:
                 table_data['target_column'] = target_column
             if predictor_id:
                 table_data['predictor_id'] = predictor_id
+            if best_metric_preference:
+                table_data['best_metric_preference'] = best_metric_preference
         
         try:
             response_data = self._post_json(f"/session/{session_id}/predict_table", table_data, max_retries=max_retries)
@@ -6136,7 +6385,8 @@ class FeatrixSphereClient:
                 raise
     
     def predict_records(self, session_id: str, records: List[Dict[str, Any]], 
-                       target_column: str = None, predictor_id: str = None, batch_size: int = 2500, use_async: bool = False, 
+                       target_column: str = None, predictor_id: str = None, best_metric_preference: str = None,
+                       batch_size: int = 2500, use_async: bool = False, 
                        show_progress_bar: bool = True, print_target_column_warning: bool = True) -> Dict[str, Any]:
         """
         Make batch predictions on a list of records with automatic client-side batching.
@@ -6183,7 +6433,8 @@ class FeatrixSphereClient:
             table_data = JSONTablesEncoder.from_records(cleaned_records)
             
             try:
-                result = self.predict_table(session_id, table_data)
+                result = self.predict_table(session_id, table_data, target_column=target_column, 
+                                           predictor_id=predictor_id, best_metric_preference=best_metric_preference)
                 
                 # Check if server returned an async job
                 if result.get('async') and result.get('job_id'):
@@ -6220,7 +6471,8 @@ class FeatrixSphereClient:
             table_data = JSONTablesEncoder.from_records(cleaned_records)
             
             try:
-                return self.predict_table(session_id, table_data)
+                return self.predict_table(session_id, table_data, target_column=target_column, 
+                                         predictor_id=predictor_id, best_metric_preference=best_metric_preference)
             except Exception as e:
                 if "404" in str(e) and "Single predictor not found" in str(e):
                     self._raise_predictor_not_found_error(session_id, "predict_records")
@@ -6250,7 +6502,8 @@ class FeatrixSphereClient:
                 table_data = JSONTablesEncoder.from_records(chunk_records)
                 
                 # Make prediction
-                chunk_result = self.predict_table(session_id, table_data)
+                chunk_result = self.predict_table(session_id, table_data, target_column=target_column, 
+                                                 predictor_id=predictor_id, best_metric_preference=best_metric_preference)
                 chunk_predictions = chunk_result.get('predictions', [])
                 
                 # Adjust row indices to match original dataset
@@ -6611,7 +6864,8 @@ class FeatrixSphereClient:
         print(f"\n⏰ Timeout after {max_wait_time} seconds")
         return {'status': 'timeout', 'message': f'Job did not complete within {max_wait_time} seconds'}
     
-    def predict_df(self, session_id: str, df, target_column: str = None, predictor_id: str = None, show_progress_bar: bool = True, print_target_column_warning: bool = True) -> Dict[str, Any]:
+    def predict_df(self, session_id: str, df, target_column: str = None, predictor_id: str = None, 
+                   best_metric_preference: str = None, show_progress_bar: bool = True, print_target_column_warning: bool = True) -> Dict[str, Any]:
         """
         Make batch predictions on a pandas DataFrame.
         
@@ -6636,7 +6890,8 @@ class FeatrixSphereClient:
         records = df.to_dict(orient='records')
         # Clean NaNs for JSON encoding
         cleaned_records = self.replace_nans_with_nulls(records)
-        return self.predict_records(session_id, cleaned_records, target_column=target_column, predictor_id=predictor_id, show_progress_bar=show_progress_bar, print_target_column_warning=print_target_column_warning)
+        return self.predict_records(session_id, cleaned_records, target_column=target_column, predictor_id=predictor_id, 
+                                   best_metric_preference=best_metric_preference, show_progress_bar=show_progress_bar, print_target_column_warning=print_target_column_warning)
     
     def _raise_predictor_not_found_error(self, session_id: str, method_name: str):
         """
@@ -6767,28 +7022,54 @@ class FeatrixSphereClient:
                     training_metrics = models.get('training_metrics', {})
                     if debug:
                         print(f"🔍 Debug: training_metrics available = {training_metrics.get('available')}")
+                    target_column = None
+                    metadata = {}
+                    
                     if training_metrics.get('available'):
-                        metrics_data = self.get_training_metrics(session_id)
-                        if debug:
-                            print(f"🔍 Debug: metrics_data keys = {list(metrics_data.keys())}")
-                        training_metrics_inner = metrics_data.get('training_metrics', {})
-                        if debug:
-                            print(f"🔍 Debug: training_metrics_inner keys = {list(training_metrics_inner.keys()) if training_metrics_inner else 'None'}")
-                        target_column = training_metrics_inner.get('target_column')
-                        if debug:
-                            print(f"🔍 Debug: extracted target_column = {target_column}")
-                        if target_column:
-                            # Extract metadata from training metrics
-                            metadata = self._extract_predictor_metadata(metrics_data, debug)
-                            
-                            # Generate unique predictor ID
-                            predictor_path = single_predictor.get('path', '')
+                        try:
+                            metrics_data = self.get_training_metrics(session_id)
+                            if debug:
+                                print(f"🔍 Debug: metrics_data keys = {list(metrics_data.keys())}")
+                            training_metrics_inner = metrics_data.get('training_metrics', {})
+                            if debug:
+                                print(f"🔍 Debug: training_metrics_inner keys = {list(training_metrics_inner.keys()) if training_metrics_inner else 'None'}")
+                            target_column = training_metrics_inner.get('target_column')
+                            if debug:
+                                print(f"🔍 Debug: extracted target_column = {target_column}")
+                            if target_column:
+                                # Extract metadata from training metrics
+                                metadata = self._extract_predictor_metadata(metrics_data, debug)
+                        except Exception as e:
+                            if debug:
+                                print(f"⚠️ Could not get training metrics: {e}")
+                    
+                    # Fallback: try to get target column from job_plan
+                    if not target_column:
+                        job_plan = session.get('job_plan', [])
+                        for job in job_plan:
+                            if job.get('job_type') == 'train_single_predictor':
+                                spec = job.get('spec', {})
+                                target_column = spec.get('target_column')
+                                if target_column:
+                                    if debug:
+                                        print(f"🔍 Debug: extracted target_column from job_plan: {target_column}")
+                                    break
+                    
+                    # If predictor is available, add it even without target_column (can be None)
+                    if single_predictor.get('available') or single_predictor.get('predictors'):
+                        # Generate unique predictor ID
+                        predictor_path = single_predictor.get('path', '')
+                        if not predictor_path and single_predictor.get('predictors'):
+                            # Use first predictor from new format
+                            predictor_path = single_predictor.get('predictors', [{}])[0].get('path', '')
+                        
+                        if predictor_path:
                             predictor_id = self._generate_predictor_id(predictor_path, 'single_predictor')
                             
                             predictors[predictor_id] = {
                                 'predictor_id': predictor_id,
                                 'path': predictor_path,
-                                'target_column': target_column,
+                                'target_column': target_column,  # Can be None
                                 'available': True,
                                 'type': 'single_predictor',
                                 **metadata  # Include epochs, validation_loss, job_status, etc.
@@ -6811,6 +7092,9 @@ class FeatrixSphereClient:
                 if debug:
                     print(f"🔍 Debug: single_predictors array = {single_predictors_paths}")
                 if single_predictors_paths:
+                    target_column = None
+                    metadata = {}
+                    
                     # Try to get target column info from training metrics
                     training_metrics = models.get('training_metrics', {})
                     if training_metrics.get('available'):
@@ -6820,30 +7104,44 @@ class FeatrixSphereClient:
                             if target_column:
                                 # Extract metadata from training metrics
                                 metadata = self._extract_predictor_metadata(metrics_data, debug)
-                                
-                                # Add each predictor individually with its own predictor_id key
-                                for i, path in enumerate(single_predictors_paths):
-                                    predictor_id = self._generate_predictor_id(path, f'multiple_predictor_{i}')
-                                    
-                                    predictors[predictor_id] = {
-                                        'predictor_id': predictor_id,
-                                        'path': path,
-                                        'target_column': target_column,
-                                        'available': True,
-                                        'type': 'single_predictor',  # Each is treated as individual predictor
-                                        'predictor_index': i,  # Track original index for compatibility
-                                        **metadata  # Include epochs, validation_loss, job_status, etc.
-                                    }
-                                    if debug:
-                                        print(f"✅ Added predictor {i} for target_column: {target_column}")
-                                        print(f"   Predictor ID: {predictor_id}")
-                                        print(f"   Path: {path}")
-                                
-                                if debug:
-                                    print(f"   Total predictors added: {len(single_predictors_paths)}")
-                                    print(f"   Shared metadata: {metadata}")
                         except Exception as e:
-                            print(f"Warning: Could not extract target column from training metrics: {e}")
+                            if debug:
+                                print(f"⚠️ Could not get training metrics: {e}")
+                    
+                    # Fallback: try to get target column from job_plan
+                    if not target_column:
+                        job_plan = session.get('job_plan', [])
+                        for job in job_plan:
+                            if job.get('job_type') == 'train_single_predictor':
+                                spec = job.get('spec', {})
+                                target_column = spec.get('target_column')
+                                if target_column:
+                                    if debug:
+                                        print(f"🔍 Debug: extracted target_column from job_plan: {target_column}")
+                                    break
+                    
+                    # Add each predictor even if target_column is None
+                    for i, path in enumerate(single_predictors_paths):
+                        predictor_id = self._generate_predictor_id(path, f'multiple_predictor_{i}')
+                        
+                        predictors[predictor_id] = {
+                            'predictor_id': predictor_id,
+                            'path': path,
+                            'target_column': target_column,  # Can be None
+                            'available': True,
+                            'type': 'single_predictor',  # Each is treated as individual predictor
+                            'predictor_index': i,  # Track original index for compatibility
+                            **metadata  # Include epochs, validation_loss, job_status, etc.
+                        }
+                        if debug:
+                            print(f"✅ Added predictor {i} for target_column: {target_column}")
+                            print(f"   Predictor ID: {predictor_id}")
+                            print(f"   Path: {path}")
+                    
+                    if debug:
+                        print(f"   Total predictors added: {len(single_predictors_paths)}")
+                        if metadata:
+                            print(f"   Shared metadata: {metadata}")
                 
                 # Fallback: check old format single_predictor field
                 single_predictor_path = session.get('single_predictor')
@@ -6920,7 +7218,7 @@ class FeatrixSphereClient:
             target_column: Specific target column to validate, or None for auto-detect
             
         Returns:
-            Validated target column name
+            Validated target column name (or None if can't determine, server will handle)
             
         Raises:
             ValueError: If target_column is invalid or multiple predictors exist without specification
@@ -6928,7 +7226,8 @@ class FeatrixSphereClient:
         available_predictors = self._get_available_predictors(session_id)
         
         if not available_predictors:
-            raise ValueError(f"No trained predictors found for session {session_id}")
+            # Don't fail - let server handle discovery. Return provided target_column or None.
+            return target_column
         
         if target_column is None:
             # Auto-detect: only valid if there's exactly one predictor
@@ -6982,6 +7281,10 @@ class FeatrixSphereClient:
             # Re-raise validation errors
             raise e
         
+        # If we couldn't determine target column (server will handle), just return records as-is
+        if validated_target_column is None:
+            return records
+        
         if validated_target_column in records[0]:
             if print_warning:
                 print(f"⚠️  Warning: Removing target column '{validated_target_column}' from prediction data")

{featrixsphere-0.2.3737.dist-info → featrixsphere-0.2.4983.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: featrixsphere
-Version: 0.2.3737
+Version: 0.2.4983
 Summary: Transform any CSV into a production-ready ML model in minutes, not months.
 Home-page: https://github.com/Featrix/sphere
 Author: Featrix

featrixsphere-0.2.4983.dist-info/RECORD

@@ -0,0 +1,8 @@
+featrixsphere/__init__.py,sha256=IRi4Di6zgujpSsWTJi0VL5b8yhjUGd9z_XFXJT2uWi8,1888
+featrixsphere/client.py,sha256=kWGR7cYH0IDWNZDX5w8yPP9-2wA-3KfXxFlYl6w01wE,431295
+featrixsphere/test_client.py,sha256=4SiRbib0ms3poK0UpnUv4G0HFQSzidF3Iswo_J2cjLk,11981
+featrixsphere-0.2.4983.dist-info/METADATA,sha256=UxwPQet3orfLfAsNhHh1Bvn_QfV2ZXpuBL9AWYyjvxg,16232
+featrixsphere-0.2.4983.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+featrixsphere-0.2.4983.dist-info/entry_points.txt,sha256=QreJeYfD_VWvbEqPmMXZ3pqqlFlJ1qZb-NtqnyhEldc,51
+featrixsphere-0.2.4983.dist-info/top_level.txt,sha256=AyN4wjfzlD0hWnDieuEHX0KckphIk_aC73XCG4df5uU,14
+featrixsphere-0.2.4983.dist-info/RECORD,,

featrixsphere-0.2.3737.dist-info/RECORD

@@ -1,8 +0,0 @@
-featrixsphere/__init__.py,sha256=rPPiD3URmePTVO31B5XHkaGVl6zd-r6OIbpBxjqi9Yg,1888
-featrixsphere/client.py,sha256=XaNFHfjogpj3exISZG1Q2SIMn-NewVsUELpzN7-5I-A,416085
-featrixsphere/test_client.py,sha256=4SiRbib0ms3poK0UpnUv4G0HFQSzidF3Iswo_J2cjLk,11981
-featrixsphere-0.2.3737.dist-info/METADATA,sha256=3gUOu7cpZPdzBkw5Z_8H6StVxotvIBCkOXmOWfW951c,16232
-featrixsphere-0.2.3737.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-featrixsphere-0.2.3737.dist-info/entry_points.txt,sha256=QreJeYfD_VWvbEqPmMXZ3pqqlFlJ1qZb-NtqnyhEldc,51
-featrixsphere-0.2.3737.dist-info/top_level.txt,sha256=AyN4wjfzlD0hWnDieuEHX0KckphIk_aC73XCG4df5uU,14
-featrixsphere-0.2.3737.dist-info/RECORD,,