🚀 Deploy to HF Spaces: Model selection + Fine-tuning updates

- Add model selection (7+ transformer models)
- Add zero-shot model selection (3 NLI models)
- Improve fine-tuning with head-only and LoRA modes
- Add training run management (export/delete)
- Configure for HF Spaces (port 7860, persistent storage)
- Update database schema for model tracking
- Add comprehensive AI model presets

.gitignore

@@ -33,3 +33,6 @@ instance/
 # OS
 .DS_Store
 Thumbs.db
+
+# Models
+models/finetuned/

.hfignore

@@ -0,0 +1,75 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+venv/
+ENV/
+env/
+.venv
+
+# Environment files
+.env
+.env.*
+
+# Local development
+instance/
+*.db
+*.sqlite
+*.sqlite3
+
+# Models and cache (will be generated on HF)
+models/finetuned/*
+.cache/
+*.pth
+*.bin
+*.onnx
+
+# Git
+.git/
+.gitignore
+.gitattributes
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Documentation (keep only README.md)
+DEPLOYMENT.md
+QUICKSTART.md
+PROJECT_STRUCTURE.md
+MIGRATION_SUMMARY.md
+Claude's Plan.md
+AI_MODEL_COMPARISON.md
+TRAINING_STRATEGY.md
+ZERO_SHOT_MODEL_SELECTION.md
+HF_DEPLOYMENT_CHECKLIST.md
+
+# Test files
+test_*.py
+mock_data*.json
+
+# Local-specific files
+Dockerfile
+docker-compose.yml
+.dockerignore
+gunicorn_config.py
+run.py
+start.sh
+
+# Keep these for HF:
+# - Dockerfile (will be copied from Dockerfile.hf)
+# - README.md (will be copied from README_HF.md)
+# - app_hf.py
+# - wsgi.py
+# - requirements.txt
+# - app/ directory
+

app/analyzer.py

@@ -90,7 +90,8 @@ class SubmissionAnalyzer:
                     finetuned_path,
                     num_labels=len(self.categories),
                     id2label=self.id2label,
-                    label2id=self.label2id
+                    label2id=self.label2id,
+                    ignore_mismatched_sizes=True
                 )
                 self.model.eval()
                 self.model_type = 'finetuned'
@@ -102,14 +103,23 @@ class SubmissionAnalyzer:
 
         # Load base zero-shot model
         try:
-            logger.info("Loading base zero-shot classification model...")
+            # Get selected zero-shot model from settings
+            from app.models.models import Settings
+            from app.fine_tuning.model_presets import get_model_preset
+            
+            zero_shot_model_key = Settings.get_setting('zero_shot_model', 'bart-large-mnli')
+            model_preset = get_model_preset(zero_shot_model_key)
+            zero_shot_model_id = model_preset['model_id']
+            
+            logger.info(f"Loading zero-shot classification model: {zero_shot_model_id}...")
             self.classifier = pipeline(
                 "zero-shot-classification",
-                model="facebook/bart-large-mnli",
+                model=zero_shot_model_id,
                 device=-1  # Use CPU (-1), change to 0 for GPU
             )
             self.model_type = 'base'
-            logger.info("Base model loaded successfully!")
+            self.zero_shot_model_key = zero_shot_model_key
+            logger.info(f"Zero-shot model loaded successfully: {model_preset['name']}!")
         except Exception as e:
             logger.error(f"Error loading model: {e}")
             raise

app/fine_tuning/model_manager.py

@@ -20,16 +20,27 @@ logger = logging.getLogger(__name__)
 class ModelManager:
     """Manage fine-tuned model deployment and versioning"""
 
-    def __init__(self, models_dir: str = "/data/models/finetuned"):
+    def __init__(self, models_dir: str = None):
         """
         Initialize ModelManager.
 
         Args:
             models_dir: Base directory for storing fine-tuned models
+                       (defaults to MODELS_DIR env var or './models/finetuned')
         """
+        if models_dir is None:
+            # Use environment variable or local path
+            models_dir = os.getenv('MODELS_DIR', 'models/finetuned')
+        
         self.models_dir = models_dir
         self.base_model_name = "facebook/bart-large-mnli"
-        os.makedirs(models_dir, exist_ok=True)
+        
+        # Create directory if it doesn't exist
+        try:
+            os.makedirs(models_dir, exist_ok=True)
+        except PermissionError:
+            logger.error(f"Permission denied creating models directory: {models_dir}")
+            raise
 
     def get_model_path(self, run_id: int) -> str:
         """Get path to model for a specific training run"""
@@ -56,7 +67,10 @@ class ModelManager:
             model_name = model_path
 
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            ignore_mismatched_sizes=True
+        )
 
         return model, tokenizer
 

app/fine_tuning/model_presets.py

@@ -0,0 +1,168 @@
+"""
+Model presets for both fine-tuning and zero-shot classification.
+Provides configuration for various HuggingFace models optimized for text classification.
+"""
+
+MODEL_PRESETS = {
+    # Zero-shot capable models (NLI-trained)
+    'bart-large-mnli': {
+        'name': 'BART-large-MNLI',
+        'model_id': 'facebook/bart-large-mnli',
+        'max_length': 1024,
+        'size': '400M',
+        'speed': 'Slow',
+        'best_for': 'Zero-shot + Fine-tuning',
+        'description': 'Large sequence-to-sequence model, excellent zero-shot performance',
+        'recommended_lr': 2e-5,
+        'recommended_batch': 4,
+        'supports_zero_shot': True
+    },
+    'deberta-v3-base-mnli': {
+        'name': 'DeBERTa-v3-base-MNLI',
+        'model_id': 'MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli',
+        'max_length': 512,
+        'size': '86M',
+        'speed': 'Fast',
+        'best_for': 'Fast zero-shot classification',
+        'description': 'DeBERTa trained on NLI datasets, excellent zero-shot with better speed',
+        'recommended_lr': 2e-5,
+        'recommended_batch': 8,
+        'supports_zero_shot': True
+    },
+    'distilbart-mnli': {
+        'name': 'DistilBART-MNLI',
+        'model_id': 'valhalla/distilbart-mnli-12-3',
+        'max_length': 1024,
+        'size': '134M',
+        'speed': 'Medium',
+        'best_for': 'Balanced zero-shot',
+        'description': 'Distilled BART for zero-shot, good balance of speed and accuracy',
+        'recommended_lr': 2e-5,
+        'recommended_batch': 8,
+        'supports_zero_shot': True
+    },
+    
+    # Fine-tuning only models
+    'deberta-v3-small': {
+        'name': 'DeBERTa-v3-small',
+        'model_id': 'microsoft/deberta-v3-small',
+        'max_length': 512,
+        'size': '44M',
+        'speed': 'Very Fast',
+        'best_for': 'Fine-tuning with small datasets',
+        'description': 'State-of-the-art efficient model, excellent for small datasets',
+        'recommended_lr': 3e-5,
+        'recommended_batch': 8,
+        'supports_zero_shot': False
+    },
+    'deberta-v3-base': {
+        'name': 'DeBERTa-v3-base',
+        'model_id': 'microsoft/deberta-v3-base',
+        'max_length': 512,
+        'size': '86M',
+        'speed': 'Fast',
+        'best_for': 'High accuracy fine-tuning',
+        'description': 'Larger DeBERTa model with better accuracy',
+        'recommended_lr': 2e-5,
+        'recommended_batch': 8,
+        'supports_zero_shot': False
+    },
+    'distilbert-base': {
+        'name': 'DistilBERT-base',
+        'model_id': 'distilbert-base-uncased',
+        'max_length': 512,
+        'size': '66M',
+        'speed': 'Fast',
+        'best_for': 'Balanced speed and accuracy',
+        'description': 'Distilled BERT, 60% faster with 97% performance retention',
+        'recommended_lr': 5e-5,
+        'recommended_batch': 8,
+        'supports_zero_shot': False
+    },
+    'roberta-base': {
+        'name': 'RoBERTa-base',
+        'model_id': 'roberta-base',
+        'max_length': 512,
+        'size': '125M',
+        'speed': 'Medium',
+        'best_for': 'Maximum accuracy',
+        'description': 'Robustly optimized BERT, excellent classification performance',
+        'recommended_lr': 2e-5,
+        'recommended_batch': 8,
+        'supports_zero_shot': False
+    },
+    'electra-small': {
+        'name': 'ELECTRA-small',
+        'model_id': 'google/electra-small-discriminator',
+        'max_length': 512,
+        'size': '14M',
+        'speed': 'Fastest',
+        'best_for': 'Speed-critical applications',
+        'description': 'Very fast and lightweight, good for production',
+        'recommended_lr': 5e-5,
+        'recommended_batch': 16,
+        'supports_zero_shot': False
+    },
+    'minilm': {
+        'name': 'MiniLM-L12',
+        'model_id': 'microsoft/MiniLM-L12-H384-uncased',
+        'max_length': 512,
+        'size': '33M',
+        'speed': 'Very Fast',
+        'best_for': 'Lightweight production deployment',
+        'description': 'Compact model optimized for speed',
+        'recommended_lr': 4e-5,
+        'recommended_batch': 12,
+        'supports_zero_shot': False
+    }
+}
+
+def get_model_preset(preset_key):
+    """Get model preset configuration by key."""
+    return MODEL_PRESETS.get(preset_key, MODEL_PRESETS['bart-large-mnli'])
+
+def get_available_models():
+    """Get list of all available models for selection."""
+    return [
+        {
+            'key': key,
+            'name': config['name'],
+            'size': config['size'],
+            'speed': config['speed'],
+            'best_for': config['best_for'],
+            'supports_zero_shot': config['supports_zero_shot']
+        }
+        for key, config in MODEL_PRESETS.items()
+    ]
+
+def get_zero_shot_models():
+    """Get list of models that support zero-shot classification."""
+    return [
+        {
+            'key': key,
+            'name': config['name'],
+            'model_id': config['model_id'],
+            'size': config['size'],
+            'speed': config['speed'],
+            'description': config['description']
+        }
+        for key, config in MODEL_PRESETS.items()
+        if config.get('supports_zero_shot', False)
+    ]
+
+def get_recommended_hyperparams(preset_key, training_mode='lora'):
+    """Get recommended hyperparameters for a model preset."""
+    preset = get_model_preset(preset_key)
+    
+    base_params = {
+        'learning_rate': preset['recommended_lr'],
+        'batch_size': preset['recommended_batch'],
+        'max_length': preset['max_length']
+    }
+    
+    if training_mode == 'head_only':
+        # Higher learning rate for head-only training
+        base_params['learning_rate'] = preset['recommended_lr'] * 2
+    
+    return base_params
+

app/fine_tuning/trainer.py

@@ -75,12 +75,27 @@ class BARTFineTuner:
         # Validate splits
         assert abs(train_split + val_split + test_split - 1.0) < 0.01, "Splits must sum to 1.0"
 
+        num_classes = len(self.categories)
+        total_examples = len(texts)
+
+        # Calculate minimum examples needed for stratified split
+        # Need at least num_classes examples in each split
+        min_test_size = int(total_examples * test_split)
+        min_val_size = int(total_examples * val_split)
+
+        # Check if we have enough examples for stratification
+        use_stratify = (min_test_size >= num_classes and min_val_size >= num_classes)
+
+        if not use_stratify:
+            logger.warning(f"Dataset too small ({total_examples} examples) for stratified split. "
+                          f"Using random split instead.")
+
         # First split: separate test set
         train_val_texts, test_texts, train_val_labels, test_labels = train_test_split(
             texts, labels,
             test_size=test_split,
             random_state=random_state,
-            stratify=labels  # Ensure balanced splits
+            stratify=labels if use_stratify else None
         )
 
         # Second split: separate train and validation
@@ -89,7 +104,7 @@ class BARTFineTuner:
             train_val_texts, train_val_labels,
             test_size=val_size_adjusted,
             random_state=random_state,
-            stratify=train_val_labels
+            stratify=train_val_labels if use_stratify else None
         )
 
         # Tokenize datasets
@@ -126,6 +141,36 @@ class BARTFineTuner:
 
         return Dataset.from_dict(dataset_dict)
 
+    def setup_head_only_model(self) -> None:
+        """
+        Set up BART model for classification head-only fine-tuning.
+        Freezes the encoder and only trains the classification head.
+        Better for small datasets (<100 examples).
+        """
+        logger.info("Setting up BART model for head-only training")
+
+        # Load base model
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            self.base_model_name,
+            num_labels=len(self.categories),
+            id2label=self.id2label,
+            label2id=self.label2id,
+            problem_type="single_label_classification",
+            ignore_mismatched_sizes=True
+        )
+
+        # Freeze all parameters except classification head
+        for name, param in self.model.named_parameters():
+            if 'classification_head' in name or 'classifier' in name:
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+
+        # Count trainable parameters
+        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in self.model.parameters())
+        logger.info(f"Trainable params: {trainable:,} / {total:,} ({100 * trainable / total:.2f}%)")
+
     def setup_lora_model(self, lora_config: Dict) -> None:
         """
         Set up BART model with LoRA adapters.
@@ -145,7 +190,8 @@ class BARTFineTuner:
             num_labels=len(self.categories),
             id2label=self.id2label,
             label2id=self.label2id,
-            problem_type="single_label_classification"
+            problem_type="single_label_classification",
+            ignore_mismatched_sizes=True  # BART-MNLI has 3 classes, we need 6
         )
 
         # Configure LoRA
@@ -193,6 +239,10 @@ class BARTFineTuner:
         # Create output directory
         os.makedirs(output_dir, exist_ok=True)
 
+        # Force CPU training to avoid cuDNN compatibility issues on WSL2
+        use_cuda = False
+        logger.info("Using CPU for training (CUDA disabled to avoid compatibility issues)")
+
         # Training arguments
         training_args = TrainingArguments(
             output_dir=output_dir,
@@ -211,7 +261,8 @@ class BARTFineTuner:
             greater_is_better=False,
             save_total_limit=2,
             report_to="none",  # Disable wandb, tensorboard
-            fp16=torch.cuda.is_available(),  # Use mixed precision if GPU available
+            use_cpu=not use_cuda,  # Use CPU if CUDA test fails
+            fp16=use_cuda,  # Only use mixed precision with working CUDA
         )
 
         # Trainer
@@ -268,7 +319,8 @@ class BARTFineTuner:
             self.tokenizer = AutoTokenizer.from_pretrained(model_path)
             self.model = AutoModelForSequenceClassification.from_pretrained(
                 model_path,
-                num_labels=len(self.categories)
+                num_labels=len(self.categories),
+                ignore_mismatched_sizes=True
             )
 
         # Make predictions
@@ -281,11 +333,24 @@ class BARTFineTuner:
 
         with torch.no_grad():
             for i in range(len(test_dataset)):
-                batch = {k: test_dataset[i][k].unsqueeze(0).to(device) for k in ['input_ids', 'attention_mask']}
+                # Get the data - handle both tensor and list formats
+                item = test_dataset[i]
+
+                # Convert to tensors if needed
+                input_ids = torch.tensor(item['input_ids']) if isinstance(item['input_ids'], list) else item['input_ids']
+                attention_mask = torch.tensor(item['attention_mask']) if isinstance(item['attention_mask'], list) else item['attention_mask']
+                label = torch.tensor(item['labels']) if isinstance(item['labels'], list) else item['labels']
+
+                # Create batch
+                batch = {
+                    'input_ids': input_ids.unsqueeze(0).to(device),
+                    'attention_mask': attention_mask.unsqueeze(0).to(device)
+                }
+
                 outputs = self.model(**batch)
                 pred = torch.argmax(outputs.logits, dim=1).item()
                 predictions.append(pred)
-                true_labels.append(test_dataset[i]['labels'].item())
+                true_labels.append(label.item() if isinstance(label, torch.Tensor) else label)
 
         # Calculate metrics
         accuracy = accuracy_score(true_labels, predictions)

app/models/models.py

@@ -51,7 +51,7 @@ class Settings(db.Model):
 
     id = db.Column(db.Integer, primary_key=True)
     key = db.Column(db.String(50), unique=True, nullable=False)
-    value = db.Column(db.String(10), nullable=False)  # 'true' or 'false'
+    value = db.Column(db.String(100), nullable=False)  # Increased to support model IDs
 
     @staticmethod
     def get_setting(key, default='true'):

app/routes/admin.py

@@ -3,11 +3,15 @@ from app.models.models import Token, Submission, Settings, TrainingExample, Fine
 from app import db
 from app.analyzer import get_analyzer
 from functools import wraps
+from typing import Dict
 import json
 import csv
 import io
 from datetime import datetime
 import os
+import logging
+
+logger = logging.getLogger(__name__)
 
 bp = Blueprint('admin', __name__, url_prefix='/admin')
 
@@ -801,20 +805,27 @@ def _run_training_job(run_id: int, config: Dict):
                 test_split=config.get('test_split', 0.15)
             )
 
-            # Setup LoRA model
-            lora_config = {
-                'r': config.get('lora_rank', 16),
-                'lora_alpha': config.get('lora_alpha', 32),
-                'lora_dropout': config.get('lora_dropout', 0.1)
-            }
-            trainer.setup_lora_model(lora_config)
+            # Setup model based on training mode
+            training_mode = config.get('training_mode', 'head_only')
+
+            if training_mode == 'head_only':
+                # Head-only training (recommended for small datasets)
+                trainer.setup_head_only_model()
+            else:
+                # LoRA training
+                lora_config = {
+                    'r': config.get('lora_rank', 16),
+                    'lora_alpha': config.get('lora_alpha', 32),
+                    'lora_dropout': config.get('lora_dropout', 0.1)
+                }
+                trainer.setup_lora_model(lora_config)
 
             # Update status to training
             run.status = 'training'
             db.session.commit()
 
             # Train
-            models_dir = os.getenv('MODELS_DIR', '/data/models/finetuned')
+            models_dir = os.getenv('MODELS_DIR', 'models/finetuned')
             output_dir = os.path.join(models_dir, f'run_{run_id}')
 
             training_config = {
@@ -889,9 +900,14 @@ def get_training_status(run_id):
     elif run.status == 'failed':
         progress = 0
 
+    # Get training mode from config
+    config = run.get_config() if hasattr(run, 'get_config') else {}
+    training_mode = config.get('training_mode', 'lora')
+    mode_label = 'classification head only' if training_mode == 'head_only' else 'LoRA adapters'
+
     status_messages = {
         'preparing': 'Preparing training data...',
-        'training': 'Training model with LoRA...',
+        'training': f'Training model ({mode_label})...',
         'evaluating': 'Evaluating model performance...',
         'completed': 'Training completed successfully!',
         'failed': 'Training failed'
@@ -971,3 +987,372 @@ def get_run_details(run_id):
     run = FineTuningRun.query.get_or_404(run_id)
 
     return jsonify(run.to_dict())
+
+
+@bp.route('/api/set-zero-shot-model', methods=['POST'])
+@admin_required
+def set_zero_shot_model():
+    """Set the zero-shot model for classification"""
+    try:
+        from app.fine_tuning.model_presets import get_model_preset
+        from app.analyzer import reload_analyzer
+        
+        data = request.get_json()
+        model_key = data.get('model_key')
+        
+        if not model_key:
+            return jsonify({'success': False, 'error': 'No model key provided'}), 400
+        
+        # Validate model exists and supports zero-shot
+        model_preset = get_model_preset(model_key)
+        if not model_preset.get('supports_zero_shot', False):
+            return jsonify({
+                'success': False,
+                'error': 'Selected model does not support zero-shot classification'
+            }), 400
+        
+        # Save setting
+        Settings.set_setting('zero_shot_model', model_key)
+        
+        # Reload analyzer with new model
+        reload_analyzer()
+        
+        logger.info(f"Zero-shot model changed to: {model_preset['name']}")
+        
+        return jsonify({
+            'success': True,
+            'message': f"Zero-shot model changed to {model_preset['name']}",
+            'model_key': model_key,
+            'model_name': model_preset['name']
+        })
+        
+    except Exception as e:
+        logger.error(f"Error changing zero-shot model: {str(e)}")
+        return jsonify({'success': False, 'error': str(e)}), 500
+
+
+@bp.route('/api/get-zero-shot-model', methods=['GET'])
+@admin_required
+def get_zero_shot_model():
+    """Get the current zero-shot model"""
+    try:
+        from app.fine_tuning.model_presets import get_model_preset
+        
+        model_key = Settings.get_setting('zero_shot_model', 'bart-large-mnli')
+        model_preset = get_model_preset(model_key)
+        
+        return jsonify({
+            'success': True,
+            'model_key': model_key,
+            'model_name': model_preset['name'],
+            'model_info': {
+                'size': model_preset['size'],
+                'speed': model_preset['speed'],
+                'description': model_preset['description']
+            }
+        })
+        
+    except Exception as e:
+        logger.error(f"Error getting zero-shot model: {str(e)}")
+        return jsonify({'success': False, 'error': str(e)}), 500
+
+
+@bp.route('/api/delete-training-run/<int:run_id>', methods=['DELETE'])
+@admin_required
+def delete_training_run(run_id):
+    """Delete a training run and its associated files"""
+    try:
+        run = FineTuningRun.query.get_or_404(run_id)
+        
+        # Prevent deletion of active model
+        if run.is_active_model:
+            return jsonify({
+                'success': False,
+                'error': 'Cannot delete the active model. Please rollback or deploy another model first.'
+            }), 400
+        
+        # Prevent deletion of currently training runs
+        if run.status == 'training':
+            return jsonify({
+                'success': False,
+                'error': 'Cannot delete a training run that is currently in progress.'
+            }), 400
+        
+        # Delete model files if they exist
+        import shutil
+        if run.model_path and os.path.exists(run.model_path):
+            try:
+                shutil.rmtree(run.model_path)
+                logger.info(f"Deleted model files at {run.model_path}")
+            except Exception as e:
+                logger.error(f"Error deleting model files: {str(e)}")
+                # Continue with database deletion even if file deletion fails
+        
+        # Unlink training examples from this run (don't delete the examples themselves)
+        for example in run.training_examples:
+            example.training_run_id = None
+            example.used_in_training = False
+        
+        # Delete the training run from database
+        db.session.delete(run)
+        db.session.commit()
+        
+        return jsonify({
+            'success': True,
+            'message': f'Training run #{run_id} deleted successfully'
+        })
+        
+    except Exception as e:
+        db.session.rollback()
+        logger.error(f"Error deleting training run: {str(e)}")
+        return jsonify({'success': False, 'error': str(e)}), 500
+
+
+@bp.route('/api/export-model/<int:run_id>', methods=['GET'])
+@admin_required
+def export_model(run_id):
+    """Export a trained model as a downloadable ZIP file"""
+    try:
+        import tempfile
+        import shutil
+        from datetime import datetime
+        
+        run = FineTuningRun.query.get_or_404(run_id)
+        
+        if run.status != 'completed':
+            return jsonify({
+                'success': False,
+                'error': 'Can only export completed training runs'
+            }), 400
+        
+        if not run.model_path or not os.path.exists(run.model_path):
+            return jsonify({
+                'success': False,
+                'error': 'Model files not found'
+            }), 404
+        
+        # Create temporary directory for export
+        temp_dir = tempfile.mkdtemp()
+        try:
+            export_name = f"model_run_{run_id}"
+            export_path = os.path.join(temp_dir, export_name)
+            
+            # Copy model files
+            shutil.copytree(run.model_path, export_path)
+            
+            # Create model card with metadata
+            config = run.get_config()
+            results = run.get_results()
+            
+            model_card = {
+                'run_id': run_id,
+                'export_date': datetime.utcnow().isoformat(),
+                'created_at': run.created_at.isoformat() if run.created_at else None,
+                'training_mode': config.get('training_mode', 'lora'),
+                'base_model': 'facebook/bart-large-mnli',
+                'model_type': 'BART fine-tuned for text classification',
+                'task': 'Multi-class text classification',
+                'categories': ['Vision', 'Problem', 'Objectives', 'Directives', 'Values', 'Actions'],
+                'training_config': config,
+                'results': results,
+                'improvement_over_baseline': run.improvement_over_baseline,
+                'num_training_examples': run.num_training_examples,
+                'num_validation_examples': run.num_validation_examples,
+                'num_test_examples': run.num_test_examples
+            }
+            
+            with open(os.path.join(export_path, 'model_card.json'), 'w') as f:
+                json.dump(model_card, f, indent=2)
+            
+            # Create README
+            readme_content = f"""# Participatory Planning Model - Run {run_id}
+
+## Model Information
+- **Export Date**: {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}
+- **Training Mode**: {config.get('training_mode', 'lora').upper()}
+- **Base Model**: facebook/bart-large-mnli
+- **Task**: Multi-class text classification
+
+## Categories
+1. Vision
+2. Problem
+3. Objectives
+4. Directives
+5. Values
+6. Actions
+
+## Training Configuration
+- **Learning Rate**: {config.get('learning_rate', 'N/A')}
+- **Epochs**: {config.get('num_epochs', 'N/A')}
+- **Batch Size**: {config.get('batch_size', 'N/A')}
+- **Training Examples**: {run.num_training_examples}
+- **Validation Examples**: {run.num_validation_examples}
+- **Test Examples**: {run.num_test_examples}
+
+## Performance
+- **Test Accuracy**: {results.get('test_accuracy', 0)*100:.1f}%
+- **Improvement over Baseline**: {run.improvement_over_baseline*100:.1f}%
+
+## Usage
+To load this model:
+```python
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+tokenizer = AutoTokenizer.from_pretrained("./model_run_{run_id}")
+model = AutoModelForSequenceClassification.from_pretrained("./model_run_{run_id}")
+```
+
+See model_card.json for detailed metrics.
+"""
+            
+            with open(os.path.join(export_path, 'README.md'), 'w') as f:
+                f.write(readme_content)
+            
+            # Create ZIP file
+            zip_path = os.path.join(temp_dir, f"model_run_{run_id}")
+            shutil.make_archive(zip_path, 'zip', temp_dir, export_name)
+            zip_file = f"{zip_path}.zip"
+            
+            # Read ZIP file into memory before cleaning up temp dir
+            with open(zip_file, 'rb') as f:
+                zip_data = io.BytesIO(f.read())
+            
+            # Clean up temp directory
+            shutil.rmtree(temp_dir)
+            
+            # Send file from memory
+            zip_data.seek(0)
+            return send_file(
+                zip_data,
+                mimetype='application/zip',
+                as_attachment=True,
+                download_name=f'participatory_planner_model_run_{run_id}_{datetime.now().strftime("%Y%m%d")}.zip'
+            )
+        except Exception as e:
+            # Clean up temp dir if error occurs
+            if os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir)
+            raise e
+    
+    except Exception as e:
+        logger.error(f"Error exporting model: {str(e)}")
+        return jsonify({'success': False, 'error': str(e)}), 500
+
+
+@bp.route('/api/import-model', methods=['POST'])
+@admin_required
+def import_model():
+    """Import a previously exported model from ZIP file"""
+    try:
+        import tempfile
+        import zipfile
+        import shutil
+        
+        if 'file' not in request.files:
+            return jsonify({'success': False, 'error': 'No file uploaded'}), 400
+        
+        file = request.files['file']
+        
+        if file.filename == '':
+            return jsonify({'success': False, 'error': 'No file selected'}), 400
+        
+        if not file.filename.endswith('.zip'):
+            return jsonify({'success': False, 'error': 'File must be a ZIP archive'}), 400
+        
+        # Create temporary directory for extraction
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Save uploaded ZIP
+            zip_path = os.path.join(temp_dir, 'upload.zip')
+            file.save(zip_path)
+            
+            # Extract ZIP
+            extract_dir = os.path.join(temp_dir, 'extracted')
+            os.makedirs(extract_dir, exist_ok=True)
+            
+            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+                zip_ref.extractall(extract_dir)
+            
+            # Find the model directory (should be model_run_X)
+            contents = os.listdir(extract_dir)
+            if len(contents) != 1:
+                return jsonify({'success': False, 'error': 'Invalid model archive structure'}), 400
+            
+            model_dir = os.path.join(extract_dir, contents[0])
+            
+            # Validate it's a valid model
+            required_files = ['config.json']
+            model_files = ['pytorch_model.bin', 'model.safetensors']  # Either format
+            
+            has_config = os.path.exists(os.path.join(model_dir, 'config.json'))
+            has_model = any(os.path.exists(os.path.join(model_dir, f)) for f in model_files)
+            
+            if not has_config or not has_model:
+                return jsonify({
+                    'success': False,
+                    'error': 'Invalid model archive - missing required files (config.json and model weights)'
+                }), 400
+            
+            # Read model card if available
+            model_info = {}
+            model_card_path = os.path.join(model_dir, 'model_card.json')
+            if os.path.exists(model_card_path):
+                with open(model_card_path, 'r') as f:
+                    model_info = json.load(f)
+            
+            # Create new training run record
+            training_run = FineTuningRun(
+                status='completed',
+                created_at=datetime.utcnow()
+            )
+            
+            # Set config from model card if available
+            if 'training_config' in model_info:
+                training_run.set_config(model_info['training_config'])
+            else:
+                # Default config for imported models
+                training_run.set_config({
+                    'training_mode': 'imported',
+                    'imported': True,
+                    'original_filename': file.filename
+                })
+            
+            # Set metadata from model card
+            if 'num_training_examples' in model_info:
+                training_run.num_training_examples = model_info['num_training_examples']
+            if 'num_validation_examples' in model_info:
+                training_run.num_validation_examples = model_info['num_validation_examples']
+            if 'num_test_examples' in model_info:
+                training_run.num_test_examples = model_info['num_test_examples']
+            if 'results' in model_info:
+                training_run.set_results(model_info['results'])
+            if 'improvement_over_baseline' in model_info:
+                training_run.improvement_over_baseline = model_info['improvement_over_baseline']
+            
+            training_run.completed_at = datetime.utcnow()
+            
+            db.session.add(training_run)
+            db.session.commit()
+            
+            # Copy model to models directory
+            models_dir = os.getenv('MODELS_DIR', 'models/finetuned')
+            destination_path = os.path.join(models_dir, f'run_{training_run.id}')
+            
+            shutil.copytree(model_dir, destination_path)
+            training_run.model_path = destination_path
+            db.session.commit()
+            
+            logger.info(f"Model imported successfully as run {training_run.id}")
+            
+            return jsonify({
+                'success': True,
+                'run_id': training_run.id,
+                'message': f'Model imported successfully as run #{training_run.id}',
+                'model_info': model_info
+            })
+    
+    except zipfile.BadZipFile:
+        return jsonify({'success': False, 'error': 'Invalid ZIP file'}), 400
+    except Exception as e:
+        db.session.rollback()
+        logger.error(f"Error importing model: {str(e)}")
+        return jsonify({'success': False, 'error': str(e)}), 500

app/templates/admin/training.html

@@ -76,6 +76,29 @@
         {% endif %}
     </div>
     <div class="card-body">
+        <!-- Zero-Shot Model Selection Section -->
+        <div class="mb-4 pb-3 border-bottom">
+            <h6><i class="bi bi-magic"></i> Zero-Shot Classification Model</h6>
+            <p class="text-muted small">Select which model to use for classifying submissions (before fine-tuning)</p>
+            <div class="row align-items-end">
+                <div class="col-md-6">
+                    <label class="form-label">Active Model</label>
+                    <select class="form-select" id="zeroShotModelSelect" onchange="changeZeroShotModel()">
+                        <option value="bart-large-mnli">BART-large-MNLI (400M) - Current Default</option>
+                        <option value="deberta-v3-base-mnli">DeBERTa-v3-base-MNLI (86M) - Fast &amp; Accurate</option>
+                        <option value="distilbart-mnli">DistilBART-MNLI (134M) - Balanced</option>
+                    </select>
+                </div>
+                <div class="col-md-6">
+                    <div id="zeroShotModelInfo" class="alert alert-info mb-0" role="alert">
+                        <small id="zeroShotModelDescription">
+                            Loading model info...
+                        </small>
+                    </div>
+                </div>
+            </div>
+        </div>
+        
         <!-- Import Training Dataset Section -->
         <div class="mb-4">
             <h6><i class="bi bi-upload"></i> Import Training Dataset</h6>
@@ -83,7 +106,19 @@
             <div class="input-group">
                 <input type="file" class="form-control" id="trainingDatasetFile" accept=".json">
                 <button class="btn btn-outline-secondary" type="button" onclick="importTrainingDataset()">
-                    <i class="bi bi-cloud-upload"></i> Import
+                    <i class="bi bi-cloud-upload"></i> Import Dataset
+                </button>
+            </div>
+        </div>
+        
+        <!-- Import Fine-Tuned Model Section -->
+        <div class="mb-4">
+            <h6><i class="bi bi-box-arrow-in-down"></i> Import Fine-Tuned Model</h6>
+            <p class="text-muted small">Upload a previously exported model ZIP file to use it in this system</p>
+            <div class="input-group">
+                <input type="file" class="form-control" id="importModelFile" accept=".zip">
+                <button class="btn btn-outline-primary" type="button" onclick="importModel()">
+                    <i class="bi bi-download"></i> Import Model
                 </button>
             </div>
         </div>
@@ -123,23 +158,21 @@
             </div>
 
             <div class="row mb-3">
-                <div class="col-md-4">
-                    <label class="form-label">
-                        LoRA Rank
-                        <button type="button" class="btn btn-sm btn-link p-0" onclick="toggleCustomLoraRank()">
-                            <i class="bi bi-pencil-square"></i>
-                        </button>
-                    </label>
-                    <select class="form-select" id="loraRank" onchange="checkCustomLoraRank()">
-                        <option value="8">8 (Fast, less capacity)</option>
-                        <option value="16" selected>16 (Balanced)</option>
-                        <option value="32">32 (Slow, more capacity)</option>
-                        <option value="custom">Custom...</option>
+                <div class="col-md-12">
+                    <label class="form-label">Training Mode</label>
+                    <select class="form-select" id="trainingMode" onchange="updateTrainingModeUI()">
+                        <option value="head_only">Classification Head Only (Recommended for small datasets)</option>
+                        <option value="lora">LoRA Fine-Tuning (For larger datasets)</option>
                     </select>
-                    <input type="number" class="form-control mt-2" id="customLoraRank"
-                           style="display: none;" placeholder="Enter custom rank (4-64)"
-                           min="4" max="64" step="4" value="16">
+                    <p class="text-muted small mt-1">
+                        <strong>Head Only:</strong> Faster, better for &lt;100 examples. Only trains the output layer.<br>
+                        <strong>LoRA:</strong> Slower, better for &gt;100 examples. Trains adapter layers throughout the model.
+                    </p>
                 </div>
+            </div>
+
+            <!-- Common Settings (visible for both modes) -->
+            <div class="row mb-3">
                 <div class="col-md-4">
                     <label class="form-label">
                         Learning Rate
@@ -174,9 +207,6 @@
                            style="display: none;" placeholder="Enter custom epochs (1-20)"
                            min="1" max="20" value="3">
                 </div>
-            </div>
-
-            <div class="row mb-3">
                 <div class="col-md-4">
                     <label class="form-label">Batch Size</label>
                     <select class="form-select" id="batchSize">
@@ -185,6 +215,28 @@
                         <option value="16">16 (High memory)</option>
                     </select>
                 </div>
+            </div>
+
+            <!-- LoRA-specific Settings (only visible in LoRA mode) -->
+            <div id="loraSettings">
+            <div class="row mb-3">
+                <div class="col-md-4">
+                    <label class="form-label">
+                        LoRA Rank
+                        <button type="button" class="btn btn-sm btn-link p-0" onclick="toggleCustomLoraRank()">
+                            <i class="bi bi-pencil-square"></i>
+                        </button>
+                    </label>
+                    <select class="form-select" id="loraRank" onchange="checkCustomLoraRank()">
+                        <option value="8">8 (Fast, less capacity)</option>
+                        <option value="16" selected>16 (Balanced)</option>
+                        <option value="32">32 (Slow, more capacity)</option>
+                        <option value="custom">Custom...</option>
+                    </select>
+                    <input type="number" class="form-control mt-2" id="customLoraRank"
+                           style="display: none;" placeholder="Enter custom rank (4-64)"
+                           min="4" max="64" step="4" value="16">
+                </div>
                 <div class="col-md-4">
                     <label class="form-label">LoRA Alpha</label>
                     <input type="number" class="form-control" id="loraAlpha" value="32" min="8" max="128" step="8">
@@ -196,6 +248,8 @@
                     <small class="text-muted">Regularization (0.0-0.5)</small>
                 </div>
             </div>
+            </div><!-- End loraSettings -->
+
 
             <div class="d-grid gap-2">
                 <button type="button" class="btn btn-primary btn-lg" onclick="startTraining()">
@@ -283,6 +337,16 @@
                             <button class="btn btn-sm btn-info" onclick="viewRunDetails({{ run.id }})">
                                 <i class="bi bi-eye"></i> Details
                             </button>
+                            {% if run.status == 'completed' %}
+                            <a href="{{ url_for('admin.export_model', run_id=run.id) }}" class="btn btn-sm btn-success" download>
+                                <i class="bi bi-download"></i> Export
+                            </a>
+                            {% endif %}
+                            {% if not run.is_active_model and run.status != 'training' %}
+                            <button class="btn btn-sm btn-danger" onclick="deleteRun({{ run.id }})">
+                                <i class="bi bi-trash"></i> Delete
+                            </button>
+                            {% endif %}
                         </td>
                     </tr>
                     {% endfor %}
@@ -497,24 +561,158 @@ function importTrainingDataset() {
     });
 }
 
+// Import fine-tuned model function
+function importModel() {
+    const fileInput = document.getElementById('importModelFile');
+    const file = fileInput.files[0];
+
+    if (!file) {
+        alert('Please select a model ZIP file to import');
+        return;
+    }
+
+    if (!file.name.endsWith('.zip')) {
+        alert('Please select a ZIP file');
+        return;
+    }
+
+    if (!confirm('Import this fine-tuned model? It will be added to your training history and can be deployed.')) {
+        return;
+    }
+
+    const formData = new FormData();
+    formData.append('file', file);
+
+    // Show loading state
+    const button = event.target;
+    const originalText = button.innerHTML;
+    button.innerHTML = '<span class="spinner-border spinner-border-sm" role="status"></span> Importing...';
+    button.disabled = true;
+
+    fetch('{{ url_for("admin.import_model") }}', {
+        method: 'POST',
+        body: formData
+    })
+    .then(response => response.json())
+    .then(data => {
+        if (data.success) {
+            alert(`Successfully imported model as run #${data.run_id}!`);
+            location.reload();
+        } else {
+            alert('Error importing model: ' + data.error);
+            button.innerHTML = originalText;
+            button.disabled = false;
+        }
+    })
+    .catch(err => {
+        alert('Error: ' + err.message);
+        button.innerHTML = originalText;
+        button.disabled = false;
+    });
+}
+
+// Update UI based on training mode
+function updateTrainingModeUI() {
+    const mode = document.getElementById('trainingMode').value;
+    const loraSettings = document.getElementById('loraSettings');
+
+    if (mode === 'head_only') {
+        loraSettings.style.display = 'none';
+    } else {
+        loraSettings.style.display = 'block';
+    }
+}
+
+// Initialize on page load
+document.addEventListener('DOMContentLoaded', function() {
+    updateTrainingModeUI();
+    loadCurrentZeroShotModel();
+});
+
+// Load current zero-shot model
+function loadCurrentZeroShotModel() {
+    fetch('{{ url_for("admin.get_zero_shot_model") }}')
+    .then(response => response.json())
+    .then(data => {
+        if (data.success) {
+            document.getElementById('zeroShotModelSelect').value = data.model_key;
+            updateZeroShotModelDescription(data.model_info);
+        }
+    })
+    .catch(err => {
+        console.error('Error loading zero-shot model:', err);
+    });
+}
+
+// Update zero-shot model description
+function updateZeroShotModelDescription(modelInfo) {
+    const desc = document.getElementById('zeroShotModelDescription');
+    if (modelInfo) {
+        desc.innerHTML = `<strong>${modelInfo.size}</strong> parameters | Speed: <strong>${modelInfo.speed}</strong><br>${modelInfo.description}`;
+    }
+}
+
+// Change zero-shot model
+function changeZeroShotModel() {
+    const modelKey = document.getElementById('zeroShotModelSelect').value;
+    
+    if (!confirm(`Switch zero-shot model? This will reload the analyzer and may take a moment.`)) {
+        loadCurrentZeroShotModel(); // Revert selection
+        return;
+    }
+    
+    const button = event.target;
+    const originalHtml = button.parentElement.innerHTML;
+    button.disabled = true;
+    
+    fetch('{{ url_for("admin.set_zero_shot_model") }}', {
+        method: 'POST',
+        headers: {'Content-Type': 'application/json'},
+        body: JSON.stringify({model_key: modelKey})
+    })
+    .then(response => response.json())
+    .then(data => {
+        if (data.success) {
+            alert(`✓ Zero-shot model changed to ${data.model_name}!\n\nAll new classifications will use this model.`);
+            loadCurrentZeroShotModel(); // Refresh info
+        } else {
+            alert('Error changing model: ' + data.error);
+            loadCurrentZeroShotModel(); // Revert selection
+        }
+    })
+    .catch(err => {
+        alert('Error: ' + err.message);
+        loadCurrentZeroShotModel(); // Revert selection
+    })
+    .finally(() => {
+        button.disabled = false;
+    });
+}
+
 // Start training function
 function startTraining() {
     if (!confirm('Start fine-tuning the model? This will take several minutes.')) {
         return;
     }
 
+    const mode = document.getElementById('trainingMode').value;
     const config = {
         train_split: parseInt(document.getElementById('trainSplit').value) / 100,
         val_split: parseInt(document.getElementById('valSplit').value) / 100,
         test_split: parseInt(document.getElementById('testSplit').value) / 100,
-        lora_rank: getLoraRank(),
-        lora_alpha: parseInt(document.getElementById('loraAlpha').value),
-        lora_dropout: parseFloat(document.getElementById('loraDropout').value),
+        training_mode: mode,
         learning_rate: getLearningRate(),
         num_epochs: getNumEpochs(),
         batch_size: parseInt(document.getElementById('batchSize').value)
     };
 
+    // Only include LoRA settings if in LoRA mode
+    if (mode === 'lora') {
+        config.lora_rank = getLoraRank();
+        config.lora_alpha = parseInt(document.getElementById('loraAlpha').value);
+        config.lora_dropout = parseFloat(document.getElementById('loraDropout').value);
+    }
+
     // Show progress modal
     const progressModal = new bootstrap.Modal(document.getElementById('trainingProgressModal'));
     progressModal.show();
@@ -610,22 +808,67 @@ function rollbackModel() {
     });
 }
 
+// Delete training run
+function deleteRun(runId) {
+    if (!confirm('Delete this training run and all associated files? This action cannot be undone.')) {
+        return;
+    }
+
+    fetch(`{{ url_for("admin.delete_training_run", run_id=0) }}`.replace('/0', `/${runId}`), {
+        method: 'DELETE'
+    })
+    .then(response => response.json())
+    .then(data => {
+        if (data.success) {
+            alert('Training run deleted successfully');
+            location.reload();
+        } else {
+            alert('Error deleting run: ' + data.error);
+        }
+    })
+    .catch(err => {
+        alert('Error: ' + err.message);
+    });
+}
+
 // View run details
 function viewRunDetails(runId) {
     fetch(`{{ url_for("admin.get_run_details", run_id=0) }}`.replace('/0', `/${runId}`))
     .then(response => response.json())
     .then(data => {
+        const config = data.training_config || {};
+        const trainingMode = config.training_mode || 'lora';
+        const modeLabel = trainingMode === 'head_only' ? 'Classification Head Only' : 'LoRA Fine-Tuning';
+        
+        // Build configuration list based on training mode
+        let configItems = `
+            <li class="list-group-item"><strong>Mode:</strong> ${modeLabel}</li>
+            <li class="list-group-item"><strong>Learning Rate:</strong> ${config.learning_rate || 'N/A'}</li>
+            <li class="list-group-item"><strong>Epochs:</strong> ${config.num_epochs || 'N/A'}</li>
+            <li class="list-group-item"><strong>Batch Size:</strong> ${config.batch_size || 'N/A'}</li>
+        `;
+        
+        // Add LoRA-specific settings if applicable
+        if (trainingMode === 'lora') {
+            configItems += `
+                <li class="list-group-item"><strong>LoRA Rank:</strong> ${config.lora_rank || 'N/A'}</li>
+                <li class="list-group-item"><strong>LoRA Alpha:</strong> ${config.lora_alpha || 'N/A'}</li>
+                <li class="list-group-item"><strong>LoRA Dropout:</strong> ${config.lora_dropout || 'N/A'}</li>
+            `;
+        }
+        
+        configItems += `
+            <li class="list-group-item"><strong>Training Examples:</strong> ${data.num_training_examples || 'N/A'}</li>
+            <li class="list-group-item"><strong>Validation Examples:</strong> ${data.num_validation_examples || 'N/A'}</li>
+            <li class="list-group-item"><strong>Test Examples:</strong> ${data.num_test_examples || 'N/A'}</li>
+        `;
+        
         const content = `
             <div class="row">
                 <div class="col-md-6">
                     <h6>Training Configuration</h6>
                     <ul class="list-group">
-                        <li class="list-group-item"><strong>LoRA Rank:</strong> ${data.config.lora_rank}</li>
-                        <li class="list-group-item"><strong>Learning Rate:</strong> ${data.config.learning_rate}</li>
-                        <li class="list-group-item"><strong>Epochs:</strong> ${data.config.num_epochs}</li>
-                        <li class="list-group-item"><strong>Training Examples:</strong> ${data.num_training_examples}</li>
-                        <li class="list-group-item"><strong>Validation Examples:</strong> ${data.num_validation_examples}</li>
-                        <li class="list-group-item"><strong>Test Examples:</strong> ${data.num_test_examples}</li>
+                        ${configItems}
                     </ul>
                 </div>
                 <div class="col-md-6">
@@ -646,6 +889,9 @@ function viewRunDetails(runId) {
         document.getElementById('runDetailsContent').innerHTML = content;
         const modal = new bootstrap.Modal(document.getElementById('runDetailsModal'));
         modal.show();
+    })
+    .catch(err => {
+        alert('Error loading run details: ' + err.message);
     });
 }
 </script>

requirements.txt

@@ -1,13 +1,13 @@
 Flask==3.0.0
 Flask-SQLAlchemy==3.1.1
 python-dotenv==1.0.0
-transformers==4.36.0
+transformers==4.46.0
 torch==2.5.0
 sentencepiece>=0.2.0
 gunicorn==21.2.0
 
 # Fine-tuning dependencies
-peft>=0.7.0
+peft==0.13.2
 datasets>=2.14.0
 scikit-learn>=1.3.0
 matplotlib>=3.7.0