app.py

#!/usr/bin/env python3
"""
HF Spaces version of the ticker model training script.
This runs on HF's GPU infrastructure instead of local Apple M2.
"""

import json
import logging
import os
import time
from pathlib import Path
from typing import List, Tuple

import torch
from sentence_transformers import InputExample, SentenceTransformer, losses
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
from torch.utils.data import DataLoader, Dataset
from datasets import load_dataset

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class InputExampleDataset(Dataset):
    def __init__(self, examples):
        self.examples = examples
    def __len__(self):
        return len(self.examples)
    def __getitem__(self, idx):
        return self.examples[idx]


class HFSpacesTrainer:
    """GPU-optimized trainer for HF Spaces infrastructure."""
    
    def __init__(self, output_dir: str = "ticker_model_bge_v6_hf"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)
        
        # 🚀 HF SPACES GPU OPTIMIZATIONS
        self.device = self._setup_device()
        self.model = self._setup_model()
        
        # 📊 GPU-OPTIMIZED HYPERPARAMETERS
        self.batch_size = 64          # Much larger batch size with GPU
        self.eval_batch_size = 128    # Even larger for evaluation
        self.num_epochs = 3
        self.learning_rate = 2e-5
        self.warmup_steps = 1000
        self.max_grad_norm = 1.0
        
        logger.info(f"🚀 Device: {self.device}")
        logger.info(f"📊 Batch size: {self.batch_size} (GPU optimized)")
        logger.info(f"🎯 Model: BGE-base-en-v1.5")
        
    def _setup_device(self) -> torch.device:
        """Setup optimal device for HF Spaces."""
        if torch.cuda.is_available():
            device = torch.device("cuda")
            logger.info(f"🔥 Using CUDA GPU: {torch.cuda.get_device_name()}")
        else:
            device = torch.device("cpu")
            logger.info("💻 Using CPU")
        return device
        
    def _setup_model(self) -> SentenceTransformer:
        """Setup BGE-base-en-v1.5 model for GPU."""
        model_name = "BAAI/bge-base-en-v1.5"
        
        logger.info(f"🔄 Loading {model_name}...")
        model = SentenceTransformer(model_name, device=str(self.device))
        
        if hasattr(model, 'max_seq_length'):
            model.max_seq_length = 512
                
        logger.info(f"✅ Model loaded on {self.device}")
        return model
        
    def load_data_from_hub(self, dataset_name: str) -> Tuple[List[InputExample], List[InputExample]]:
        """Load training data from HF Hub."""
        logger.info(f"📂 Loading data from HF Hub: {dataset_name}")
        
        # Load dataset from HF Hub
        dataset = load_dataset(dataset_name)
        
        # Convert to InputExample format
        train_examples = []
        for item in dataset['train']:
            train_examples.append(InputExample(
                texts=[item['sentence_0'], item['sentence_1']],
                label=float(item['label'])
            ))
        
        eval_examples = []
        if 'test' in dataset:
            for item in dataset['test']:
                eval_examples.append(InputExample(
                    texts=[item['sentence_0'], item['sentence_1']],
                    label=float(item['label'])
                ))
        else:
            # Split train data for evaluation
            split_idx = int(0.8 * len(train_examples))
            eval_examples = train_examples[split_idx:]
            train_examples = train_examples[:split_idx]
        
        logger.info(f"📈 Training examples: {len(train_examples)}")
        logger.info(f"📊 Evaluation examples: {len(eval_examples)}")
        
        return train_examples, eval_examples
        
    def setup_training_components(self, train_examples: List[InputExample], eval_examples: List[InputExample]):
        """Setup loss function and evaluator."""
        # 🎯 COSINE SIMILARITY LOSS
        train_loss = losses.CosineSimilarityLoss(self.model)
        
        # 🔍 EVALUATOR
        eval_subset = eval_examples[:1000] if len(eval_examples) > 1000 else eval_examples
        eval_sentences1 = [ex.texts[0] for ex in eval_subset if ex.texts and len(ex.texts) > 0]
        eval_sentences2 = [ex.texts[1] for ex in eval_subset if ex.texts and len(ex.texts) > 1]
        eval_scores = [ex.label for ex in eval_subset if ex.texts and len(ex.texts) > 1]
        
        evaluator = EmbeddingSimilarityEvaluator(
            eval_sentences1, 
            eval_sentences2, 
            eval_scores,
            batch_size=self.eval_batch_size,
            name="ticker_resolution_eval"
        )
        
        return train_loss, evaluator
        
    def train(self, dataset_name: str):
        """Main training loop for HF Spaces."""
        logger.info("🚀 Starting HF Spaces GPU training...")
        total_start_time = time.time()
        
        # Load data from HF Hub
        train_examples, eval_examples = self.load_data_from_hub(dataset_name)
        
        # Setup training components
        train_loss, evaluator = self.setup_training_components(train_examples, eval_examples)
        
        # Create DataLoader
        train_dataset = InputExampleDataset(train_examples)
        train_dataloader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
        
        # 🏋️ TRAINING ARGUMENTS (GPU optimized)
        training_args = {
            'train_objectives': [(train_dataloader, train_loss)],
            'evaluator': evaluator,
            'epochs': self.num_epochs,
            'evaluation_steps': 1000,
            'warmup_steps': self.warmup_steps,
            'optimizer_params': {
                'lr': self.learning_rate,
                'eps': 1e-6,
                'weight_decay': 0.01
            },
            'output_path': str(self.output_dir),
            'save_best_model': True,
            'show_progress_bar': True,
            'use_amp': True,  # Enable mixed precision for GPU
        }
        
        # 🔥 START TRAINING
        logger.info(f"🎯 Training for {self.num_epochs} epochs with batch size {self.batch_size}")
        
        self.model.fit(**training_args)
        
        # 💾 SAVE FINAL MODEL
        final_model_path = self.output_dir / "final_model"
        self.model.save(str(final_model_path))
        
        # 📊 FINAL EVALUATION
        final_score = evaluator(self.model)
        if isinstance(final_score, dict):
            score_value = final_score.get('cosine_spearman', final_score.get('spearman', 0.0))
        else:
            score_value = final_score
        
        # 📈 TRAINING SUMMARY
        training_summary = {
            'total_training_time': time.time() - total_start_time,
            'training_examples': len(train_examples),
            'evaluation_examples': len(eval_examples),
            'epochs': self.num_epochs,
            'batch_size': self.batch_size,
            'learning_rate': self.learning_rate,
            'final_evaluation_score': score_value,
            'device': str(self.device),
            'model_name': "BGE-base-en-v1.5",
            'optimized_for': "HF Spaces GPU"
        }
        
        with open(self.output_dir / "training_summary.json", 'w') as f:
            json.dump(training_summary, f, indent=2)
            
        logger.info(f"✅ Training completed in {time.time() - total_start_time:.2f}s")
        logger.info(f"📊 Final evaluation score: {score_value:.4f}")
        
        # 🚀 PUSH TO HUB
        try:
            self.model.push_to_hub(f"ticker-model-bge-v6-hf-{int(time.time())}")
            logger.info("✅ Model pushed to HF Hub!")
        except Exception as e:
            logger.warning(f"Could not push to hub: {e}")
        
        return score_value


def main():
    # Environment variables should be set in HF Spaces
    dataset_name = os.environ.get("DATASET_NAME", "nicbstme/ticker_training_data_autotrain")
    
    # Create trainer
    trainer = HFSpacesTrainer()
    
    # Start training
    trainer.train(dataset_name)


if __name__ == "__main__":
    main()