train_embeddings.py

"""
Helion-V1-Embeddings Training Script
Train a lightweight embedding model for semantic similarity and retrieval
"""

import json
import logging
from typing import List, Dict, Tuple
from pathlib import Path
from datetime import datetime

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


class EmbeddingsTrainer:
    """Train embeddings model for Helion-V1-Embeddings."""
    
    def __init__(
        self,
        base_model: str = "sentence-transformers/all-MiniLM-L6-v2",
        output_path: str = "./helion-embeddings-output"
    ):
        self.base_model = base_model
        self.output_path = Path(output_path)
        self.output_path.mkdir(parents=True, exist_ok=True)
    
    def prepare_training_data(self) -> List[Dict]:
        """
        Prepare training data for embeddings.
        Format: sentence pairs with similarity scores.
        """
        training_examples = [
            # High similarity pairs
            {
                "sentence1": "How do I reset my password?",
                "sentence2": "What's the password reset process?",
                "score": 0.95
            },
            {
                "sentence1": "Machine learning training methods",
                "sentence2": "How to train ML models",
                "score": 0.90
            },
            {
                "sentence1": "Python programming tutorial",
                "sentence2": "Learn Python coding",
                "score": 0.88
            },
            
            # Medium similarity pairs
            {
                "sentence1": "Install Python on Windows",
                "sentence2": "Python setup guide",
                "score": 0.70
            },
            {
                "sentence1": "Best restaurants in Paris",
                "sentence2": "Where to eat in France",
                "score": 0.65
            },
            
            # Low similarity pairs
            {
                "sentence1": "How to bake cookies",
                "sentence2": "Machine learning algorithms",
                "score": 0.10
            },
            {
                "sentence1": "Weather forecast tomorrow",
                "sentence2": "Stock market analysis",
                "score": 0.05
            }
        ]
        
        logger.info(f"Prepared {len(training_examples)} training examples")
        return training_examples
    
    def create_contrastive_pairs(self) -> List[Tuple[str, str]]:
        """
        Create pairs for contrastive learning.
        Format: (anchor, positive) pairs.
        """
        pairs = [
            ("What is machine learning?", "Machine learning explained simply"),
            ("How to learn Python?", "Python learning resources"),
            ("Best coding practices", "Software development best practices"),
            ("Data science tutorial", "Learn data science basics"),
            ("Natural language processing", "NLP fundamentals guide"),
            ("Deep learning introduction", "Getting started with deep learning"),
            ("Web development guide", "How to build websites"),
            ("Database design principles", "SQL database design tutorial"),
            ("Cloud computing basics", "Introduction to cloud services"),
            ("API development guide", "How to create REST APIs"),
        ]
        
        logger.info(f"Created {len(pairs)} contrastive pairs")
        return pairs
    
    def train_model(
        self,
        train_examples: List[Dict] = None,
        epochs: int = 3,
        batch_size: int = 16,
        warmup_steps: int = 100
    ):
        """
        Train the embeddings model.
        
        Args:
            train_examples: Training data (if None, uses default)
            epochs: Number of training epochs
            batch_size: Batch size for training
            warmup_steps: Warmup steps for learning rate
        """
        try:
            from sentence_transformers import (
                SentenceTransformer,
                InputExample,
                losses,
                evaluation
            )
            from torch.utils.data import DataLoader
            
            logger.info("Loading base model...")
            model = SentenceTransformer(self.base_model)
            
            # Prepare data
            if train_examples is None:
                train_examples = self.prepare_training_data()
            
            # Convert to InputExample format
            train_data = []
            for example in train_examples:
                train_data.append(InputExample(
                    texts=[example["sentence1"], example["sentence2"]],
                    label=example["score"]
                ))
            
            # Create DataLoader
            train_dataloader = DataLoader(
                train_data,
                shuffle=True,
                batch_size=batch_size
            )
            
            # Define loss function
            train_loss = losses.CosineSimilarityLoss(model)
            
            # Training
            logger.info("Starting training...")
            model.fit(
                train_objectives=[(train_dataloader, train_loss)],
                epochs=epochs,
                warmup_steps=warmup_steps,
                output_path=str(self.output_path),
                show_progress_bar=True,
                save_best_model=True
            )
            
            logger.info(f"✅ Training complete! Model saved to {self.output_path}")
            
            return model
            
        except ImportError:
            logger.error("sentence-transformers not installed. Install with: pip install sentence-transformers")
            return None
        except Exception as e:
            logger.error(f"Training failed: {e}")
            return None
    
    def evaluate_model(self, model, test_pairs: List[Tuple[str, str, float]] = None):
        """
        Evaluate the trained model.
        
        Args:
            model: Trained SentenceTransformer model
            test_pairs: List of (sentence1, sentence2, expected_similarity)
        """
        from sentence_transformers import util
        
        if test_pairs is None:
            # Default test pairs
            test_pairs = [
                ("How to code?", "Coding tutorial", 0.85),
                ("Weather today", "Stock prices", 0.1),
                ("Machine learning", "AI and ML", 0.95),
            ]
        
        logger.info("Evaluating model...")
        
        total_error = 0
        for sent1, sent2, expected in test_pairs:
            emb1 = model.encode(sent1)
            emb2 = model.encode(sent2)
            similarity = float(util.cos_sim(emb1, emb2)[0][0])
            error = abs(similarity - expected)
            total_error += error
            
            logger.info(f"'{sent1}' <-> '{sent2}'")
            logger.info(f"  Expected: {expected:.2f}, Got: {similarity:.2f}, Error: {error:.2f}")
        
        avg_error = total_error / len(test_pairs)
        logger.info(f"Average error: {avg_error:.3f}")
        
        return avg_error
    
    def create_config_files(self):
        """Create necessary configuration files."""
        
        # Sentence transformers config
        config = {
            "__version__": {
                "sentence_transformers": "2.2.2",
                "transformers": "4.36.0",
                "pytorch": "2.0.0"
            },
            "prompts": {},
            "default_prompt_name": None,
            "similarity_fn_name": "cosine",
            "max_seq_length": 256,
            "do_lower_case": False
        }
        
        with open(self.output_path / "config_sentence_transformers.json", 'w') as f:
            json.dump(config, f, indent=2)
        
        # Modules configuration
        modules = [
            {
                "idx": 0,
                "name": "0",
                "path": "",
                "type": "sentence_transformers.models.Transformer"
            },
            {
                "idx": 1,
                "name": "1",
                "path": "1_Pooling",
                "type": "sentence_transformers.models.Pooling"
            },
            {
                "idx": 2,
                "name": "2",
                "path": "2_Normalize",
                "type": "sentence_transformers.models.Normalize"
            }
        ]
        
        with open(self.output_path / "modules.json", 'w') as f:
            json.dump(modules, f, indent=2)
        
        logger.info("✅ Configuration files created")


def main():
    """Main training function."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Train Helion-V1-Embeddings model"
    )
    parser.add_argument(
        "--base-model",
        default="sentence-transformers/all-MiniLM-L6-v2",
        help="Base model to fine-tune"
    )
    parser.add_argument(
        "--output",
        default="./helion-embeddings-output",
        help="Output directory"
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=3,
        help="Number of training epochs"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=16,
        help="Batch size"
    )
    parser.add_argument(
        "--data-file",
        type=str,
        help="Path to training data JSON file"
    )
    
    args = parser.parse_args()
    
    # Create trainer
    trainer = EmbeddingsTrainer(
        base_model=args.base_model,
        output_path=args.output
    )
    
    # Load custom data if provided
    train_examples = None
    if args.data_file:
        with open(args.data_file, 'r') as f:
            train_examples = json.load(f)
        logger.info(f"Loaded {len(train_examples)} examples from {args.data_file}")
    
    # Train model
    model = trainer.train_model(
        train_examples=train_examples,
        epochs=args.epochs,
        batch_size=args.batch_size
    )
    
    if model:
        # Evaluate
        trainer.evaluate_model(model)
        
        # Create config files
        trainer.create_config_files()
        
        print("\n" + "="*60)
        print("✅ Helion-V1-Embeddings Training Complete!")
        print("="*60)
        print(f"📁 Model saved to: {args.output}")
        print("\n💡 Test your model:")
        print("```python")
        print("from sentence_transformers import SentenceTransformer")
        print(f"model = SentenceTransformer('{args.output}')")
        print("embeddings = model.encode(['Hello world'])")
        print("```")
        print("="*60)


if __name__ == "__main__":
    main()