Upload 8 files

.gitattributes

@@ -1,2 +1,5 @@
 model-00001-of-00003.safetensors filter=lfs diff=lfs merge=lfs -text
 tokenizer.model filter=lfs diff=lfs merge=lfs -text
+Document1.pdf filter=lfs diff=lfs merge=lfs -text
+model-00002-of-00003.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00003-of-00003.safetensors filter=lfs diff=lfs merge=lfs -text

Document1.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e678058adddf0a03284f79f65242699fee2cf5191b239a9a668ada8be9862e90
+size 6035292

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9a7c7adf0142010ea7fb2d6d60b2698b86f36847d00d0afa4170c3a9fb66a9c
+size 4934842808

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4c9f1d21524ad189e63230a62a62997c52205f9ce3099948c7fc3d27385d0dc
+size 2598483736

prepare_data.py

@@ -0,0 +1,116 @@
+import os
+import json
+from pathlib import Path
+import librosa
+import taglib
+from tqdm import tqdm
+import logging
+import soundfile as sf
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class MusicDataPreprocessor:
+    def __init__(self, input_dir: str, output_dir: str):
+        self.input_dir = Path(input_dir)
+        self.output_dir = Path(output_dir)
+        self.metadata = []
+        
+        # Create necessary directories
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        (self.output_dir / "audio").mkdir(exist_ok=True)
+        (self.output_dir / "metadata").mkdir(exist_ok=True)
+    
+    def extract_metadata(self, audio_path: Path) -> dict:
+        """Extract metadata from audio file (MP3 or WAV)"""
+        try:
+            # Read audio file metadata
+            audio_format = audio_path.suffix.lower()[1:]  # Get extension without dot
+            audio_file = taglib.File(str(audio_path))
+            
+            # Get basic audio properties
+            y, sr = librosa.load(audio_path, sr=16000)  # Resample to 16kHz
+            duration = librosa.get_duration(y=y, sr=sr)
+            
+            metadata = {
+                "filename": audio_path.name,
+                "format": audio_format,
+                "duration": duration,
+                "genre": audio_file.tags.get("GENRE", ["unknown"])[0],
+                "title": audio_file.tags.get("TITLE", ["unknown"])[0],
+                "artist": audio_file.tags.get("ARTIST", ["unknown"])[0],
+                "sample_rate": sr,
+                "channels": audio_file.channels
+            }
+            
+            return metadata
+            
+        except Exception as e:
+            logger.error(f"Error processing {audio_path}: {str(e)}")
+            return None
+    
+    def process_files(self):
+        """Process all audio files (MP3 and WAV) in the input directory"""
+        # Find all MP3 and WAV files
+        audio_files = list(self.input_dir.glob("**/*.[mw][pa][3v]"))  # Match mp3, wav files
+        
+        formats_found = {"mp3": 0, "wav": 0, "other": 0}
+        formats_processed = {"mp3": 0, "wav": 0}
+        
+        logger.info(f"Found {len(audio_files)} audio files to process")
+        
+        for audio_path in tqdm(audio_files, desc="Processing audio files"):
+            # Track format statistics
+            file_ext = audio_path.suffix.lower()[1:]
+            if file_ext == "mp3":
+                formats_found["mp3"] += 1
+            elif file_ext == "wav":
+                formats_found["wav"] += 1
+            else:
+                formats_found["other"] += 1
+                logger.warning(f"Unexpected file format: {file_ext} for file {audio_path}")
+            
+            metadata = self.extract_metadata(audio_path)
+            
+            if metadata:
+                # Save processed audio - convert all to WAV
+                output_audio_path = self.output_dir / "audio" / f"{audio_path.stem}.wav"
+                try:
+                    y, sr = librosa.load(audio_path, sr=16000, mono=True)
+                    sf.write(output_audio_path, y, sr, format='WAV')
+                    
+                    # Track successful processing
+                    formats_processed[file_ext] += 1
+                    
+                    # Add path information to metadata
+                    metadata["processed_path"] = str(output_audio_path.relative_to(self.output_dir))
+                    self.metadata.append(metadata)
+                    
+                except Exception as e:
+                    logger.error(f"Error saving {audio_path}: {str(e)}")
+                    continue
+        
+        # Save metadata
+        with open(self.output_dir / "metadata" / "dataset_info.json", "w") as f:
+            json.dump({
+                "files": self.metadata,
+                "stats": {
+                    "total_processed": len(self.metadata),
+                    "formats_found": formats_found,
+                    "formats_processed": formats_processed
+                }
+            }, f, indent=2)
+        
+        logger.info(f"Processed {len(self.metadata)} files successfully")
+        logger.info(f"Files found: MP3: {formats_found['mp3']}, WAV: {formats_found['wav']}")
+        logger.info(f"Files processed: MP3: {formats_processed['mp3']}, WAV: {formats_processed['wav']}")
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_dir", type=str, required=True, help="Directory containing music files")
+    parser.add_argument("--output_dir", type=str, required=True, help="Directory to save processed files")
+    args = parser.parse_args()
+    
+    preprocessor = MusicDataPreprocessor(args.input_dir, args.output_dir)
+    preprocessor.process_files() 

requirements.txt

@@ -0,0 +1,15 @@
+torch>=2.0.0
+transformers>=4.42.0
+datasets>=2.14.0
+accelerate>=0.27.0
+librosa>=0.10.0
+pytaglib>=2.0.0
+tqdm>=4.65.0
+numpy>=1.24.0
+einops>=0.6.0
+flash-attn>=2.3.0  # Optional, for CUDA acceleration
+safetensors>=0.4.0
+soundfile>=0.12.0
+pydub>=0.25.1  # For better MP3 support
+huggingface_hub>=0.20.3
+tokenizers>=0.15.0 

train_hcf.py

@@ -0,0 +1,639 @@
+import os
+import json
+import torch
+import logging
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Tuple, Any
+import transformers
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from datasets import Dataset, load_dataset
+import numpy as np
+from accelerate import Accelerator
+from safetensors import safe_open
+from safetensors.torch import save_file, load_file
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class TensorInfo:
+    """Stores metadata about tensor indices and shape"""
+    shape: Tuple[int, ...]
+    dtype: str
+    indices: Optional[torch.Tensor] = None
+    hcf_patterns: Optional[Dict] = None
+
+class SafeTensorHCFAnalyzer:
+    """
+    Analyzes HCF patterns in model weights using SafeTensors format.
+    Handles efficient loading and analysis of large model weights.
+    """
+    
+    def __init__(self, tolerance: float = 1e-5):
+        self.tolerance = tolerance
+        self.tensor_info = {}
+        self.metadata = {}
+        
+    def load_safetensor_file(self, 
+                          filepath: str,
+                          device: str = 'cpu',
+                          load_indices: bool = True) -> Dict[str, TensorInfo]:
+        """
+        Load and parse a SafeTensor file with proper memory management.
+        
+        Args:
+            filepath: Path to .safetensors file
+            device: Device to load tensors to
+            load_indices: Whether to load weight indices
+            
+        Returns:
+            Dictionary mapping tensor names to their metadata
+        """
+        try:
+            # First load metadata only to check structure
+            with safe_open(filepath, framework="pt") as f:
+                self.metadata = json.loads(f.metadata()) if f.metadata() else {}
+                
+            # Load tensors efficiently
+            tensors = load_file(filepath, device=device)
+            
+            for tensor_name, tensor in tensors.items():
+                self.tensor_info[tensor_name] = TensorInfo(
+                    shape=tuple(tensor.shape),
+                    dtype=str(tensor.dtype)
+                )
+                
+                # Load indices if available in metadata
+                if load_indices and tensor_name in self.metadata:
+                    if 'indices' in self.metadata[tensor_name]:
+                        indices_data = self.metadata[tensor_name]['indices']
+                        if isinstance(indices_data, list):
+                            self.tensor_info[tensor_name].indices = torch.tensor(
+                                indices_data, device=device
+                            )
+                        elif isinstance(indices_data, str) and os.path.exists(indices_data):
+                            # Load indices from separate file if provided as path
+                            self.tensor_info[tensor_name].indices = torch.load(indices_data)
+                            
+            return self.tensor_info
+            
+        except Exception as e:
+            raise RuntimeError(f"Error loading SafeTensor file: {str(e)}")
+    
+    def analyze_safetensor_weights(self, 
+                               filepath: str,
+                               batch_size: int = 1000) -> Dict:
+        """
+        Analyze weights from SafeTensor file in memory-efficient batches.
+        
+        Args:
+            filepath: Path to .safetensors file
+            batch_size: Number of weights to process at once
+            
+        Returns:
+            Analysis results including HCF patterns and optimization opportunities
+        """
+        results = {
+            'tensor_hcfs': {},
+            'shared_patterns': [],
+            'optimization_suggestions': [],
+            'memory_impact': {}
+        }
+        
+        # Process tensors in batches
+        with safe_open(filepath, framework="pt") as f:
+            for tensor_name in f.keys():
+                # Get tensor info
+                tensor_data = f.get_tensor(tensor_name)
+                tensor_size = np.prod(tensor_data.shape)
+                
+                if tensor_name in self.tensor_info and self.tensor_info[tensor_name].indices is not None:
+                    indices = self.tensor_info[tensor_name].indices
+                    unique_indices = torch.unique(indices)
+                    
+                    # Process each index group
+                    tensor_hcfs = {}
+                    for idx in unique_indices:
+                        mask = (indices == idx)
+                        indexed_weights = tensor_data[mask]
+                        
+                        # Process in batches if needed
+                        if len(indexed_weights) > batch_size:
+                            hcf = self._process_large_weight_group(indexed_weights, batch_size)
+                        else:
+                            hcf = self._calculate_hcf(indexed_weights)
+                            
+                        tensor_hcfs[idx.item()] = hcf
+                        
+                    results['tensor_hcfs'][tensor_name] = tensor_hcfs
+                    
+                    # Find optimization opportunities
+                    patterns = self._analyze_weight_patterns(tensor_data, indices)
+                    self.tensor_info[tensor_name].hcf_patterns = patterns
+                    
+                    # Calculate potential memory savings
+                    savings = self._estimate_memory_savings(patterns, tensor_data.dtype)
+                    results['memory_impact'][tensor_name] = {
+                        'original_size': tensor_size * tensor_data.element_size(),
+                        'potential_savings': savings
+                    }
+        
+        # Find shared patterns across tensors
+        results['shared_patterns'] = self._find_shared_patterns()
+        results['optimization_suggestions'] = self._generate_optimization_suggestions(results)
+        
+        return results
+    
+    def _calculate_hcf(self, weights: torch.Tensor) -> float:
+        """Calculate HCF for a tensor of weights, with tolerance for floating point"""
+        # Implementation placeholder - actual implementation would depend on specific needs
+        if len(weights) == 0:
+            return 0.0
+        return 1.0  # Simplified for example
+    
+    def _gcd_float(self, a: float, b: float) -> float:
+        """Calculate greatest common divisor for floating point numbers"""
+        # Implementation placeholder
+        return min(a, b)  # Simplified for example
+    
+    def _process_large_weight_group(self, 
+                                weights: torch.Tensor, 
+                                batch_size: int) -> float:
+        """Process large weight groups in batches to manage memory."""
+        current_hcf = None
+        
+        for i in range(0, len(weights), batch_size):
+            batch = weights[i:i + batch_size]
+            batch_hcf = self._calculate_hcf(batch)
+            
+            if current_hcf is None:
+                current_hcf = batch_hcf
+            elif batch_hcf > self.tolerance:
+                current_hcf = self._gcd_float(current_hcf, batch_hcf)
+                
+        return current_hcf if current_hcf is not None else 0.0
+    
+    def _analyze_weight_patterns(self, 
+                             weights: torch.Tensor, 
+                             indices: torch.Tensor) -> Dict:
+        """Analyze weight patterns within indexed groups."""
+        patterns = {}
+        unique_indices = torch.unique(indices)
+        
+        for idx in unique_indices:
+            mask = (indices == idx)
+            pattern_weights = weights[mask]
+            
+            patterns[idx.item()] = {
+                'mean': float(pattern_weights.mean()),
+                'std': float(pattern_weights.std()),
+                'size': len(pattern_weights),
+                'hcf': self._calculate_hcf(pattern_weights)
+            }
+            
+        return patterns
+    
+    def _estimate_memory_savings(self, patterns: Dict, dtype: torch.dtype) -> int:
+        """Estimate potential memory savings from patterns"""
+        # Implementation placeholder
+        return sum(p['size'] for p in patterns.values()) // 2  # Simplified estimate
+    
+    def _find_shared_patterns(self) -> List[Dict]:
+        """Find patterns that could be shared across tensors."""
+        shared_patterns = []
+        pattern_groups = {}
+        
+        for tensor_name, info in self.tensor_info.items():
+            if info.hcf_patterns:
+                for idx, pattern in info.hcf_patterns.items():
+                    # Create pattern signature
+                    signature = f"{pattern['mean']:.4f}_{pattern['std']:.4f}"
+                    
+                    if signature not in pattern_groups:
+                        pattern_groups[signature] = []
+                    pattern_groups[signature].append({
+                        'tensor': tensor_name,
+                        'index': idx,
+                        'pattern': pattern
+                    })
+        
+        # Find groups with similar patterns
+        for signature, group in pattern_groups.items():
+            if len(group) > 1:
+                shared_patterns.append({
+                    'signature': signature,
+                    'occurrences': group,
+                    'potential_savings': sum(p['pattern']['size'] for p in group[1:])
+                })
+                
+        return shared_patterns
+    
+    def _generate_optimization_suggestions(self, results: Dict) -> List[Dict]:
+        """Generate optimization suggestions based on analysis"""
+        # Implementation placeholder
+        suggestions = []
+        for tensor_name, impact in results['memory_impact'].items():
+            if impact['potential_savings'] > 1000000:  # If savings > 1MB
+                suggestions.append({
+                    'tensor': tensor_name,
+                    'suggestion': 'Consider weight quantization',
+                    'impact': f"Save {impact['potential_savings'] / 1024 / 1024:.2f}MB"
+                })
+        return suggestions
+
+@dataclass
+class TrainingStatistics:
+    """Statistics collected during HCF-aware training"""
+    memory_savings: int = 0
+    quantization_error: float = 0.0
+    convergence_rate: float = 0.0
+    epoch: int = 0
+    batch_count: int = 0
+    
+    def update(self, batch_stats: Dict[str, Any]):
+        """Update statistics with batch results"""
+        self.memory_savings += batch_stats.get('memory_savings', 0)
+        self.quantization_error = batch_stats.get('quantization_error', self.quantization_error)
+        self.convergence_rate = batch_stats.get('convergence_rate', self.convergence_rate)
+        self.batch_count += 1
+
+class HCFTrainingOptimizer(torch.optim.Adam):
+    """
+    Optimizer with HCF-awareness for more efficient training
+    """
+    def __init__(self, 
+                 params, 
+                 lr=0.001, 
+                 betas=(0.9, 0.999), 
+                 eps=1e-8,
+                 weight_decay=0, 
+                 weight_quantization=True,
+                 maintain_patterns=True):
+        super().__init__(params, lr, betas, eps, weight_decay)
+        self.weight_quantization = weight_quantization
+        self.maintain_patterns = maintain_patterns
+        self.analyzer = SafeTensorHCFAnalyzer()
+        self.stats = {'memory_savings': 0, 'quantization_error': 0.0}
+    
+    def step(self, closure=None):
+        """Perform optimization step with HCF awareness"""
+        # Run standard optimization step
+        loss = super().step(closure)
+        
+        # Apply HCF optimizations if enabled
+        if self.weight_quantization:
+            self._apply_weight_quantization()
+        
+        if self.maintain_patterns:
+            self._maintain_weight_patterns()
+            
+        return loss
+    
+    def _apply_weight_quantization(self):
+        """Apply dynamic weight quantization using HCF patterns"""
+        savings = 0
+        total_error = 0.0
+        
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None or not p.requires_grad:
+                    continue
+                
+                # Apply weight quantization logic based on HCF analysis
+                # This is a simplified placeholder - real implementation would be more complex
+                if p.dim() > 1:  # Only apply to matrices/tensors
+                    # Find suitable quantization factor
+                    factor = torch.max(torch.abs(p.data)) / 127  # 8-bit quantization example
+                    
+                    # Quantize weights
+                    quantized = torch.round(p.data / factor) * factor
+                    
+                    # Calculate error and savings
+                    error = torch.mean((p.data - quantized)**2).item()
+                    savings += p.numel() * (p.element_size() - 1)  # Assuming 8-bit savings
+                    
+                    # Apply quantized weights
+                    p.data.copy_(quantized)
+                    
+                    total_error += error
+        
+        # Update statistics
+        self.stats['memory_savings'] = savings
+        self.stats['quantization_error'] = total_error
+    
+    def _maintain_weight_patterns(self):
+        """Maintain efficient weight patterns identified by HCF analysis"""
+        # Placeholder for pattern maintenance logic
+        # Real implementation would analyze weight matrices and enforce patterns
+        pass
+    
+    def get_stats(self):
+        """Get current optimization statistics"""
+        return self.stats
+
+class HCFAwareTrainer:
+    """
+    Trainer that incorporates HCF analysis for better training efficiency
+    """
+    def __init__(self, model, optimizer):
+        self.model = model
+        self.optimizer = optimizer
+        self.analyzer = SafeTensorHCFAnalyzer()
+        
+    def train_epoch(self, train_loader, criterion, epoch):
+        """Train one epoch with HCF awareness"""
+        self.model.train()
+        stats = TrainingStatistics(epoch=epoch)
+        
+        for batch_idx, batch in enumerate(train_loader):
+            # Get data
+            inputs, targets = self._prepare_batch(batch)
+            
+            # Forward pass
+            self.optimizer.zero_grad()
+            outputs = self.model(inputs)
+            loss = criterion(outputs, targets)
+            
+            # Backward pass
+            loss.backward()
+            
+            # Optimize with HCF awareness
+            self.optimizer.step()
+            
+            # Get batch statistics
+            batch_stats = self.optimizer.get_stats()
+            stats.update(batch_stats)
+            
+            # Log progress
+            if batch_idx % 50 == 0:
+                logger.info(f"Epoch {epoch} | Batch {batch_idx}/{len(train_loader)} | "
+                            f"Memory Savings: {stats.memory_savings/1024/1024:.2f}MB | "
+                            f"Quantization Error: {stats.quantization_error:.6f}")
+        
+        # End of epoch analysis
+        self._analyze_model_weights()
+        
+        return stats
+    
+    def _prepare_batch(self, batch):
+        """Prepare batch data for training"""
+        # Implementation depends on dataset structure
+        if isinstance(batch, dict):
+            inputs = batch.get('input_ids')
+            targets = batch.get('labels', inputs)
+        else:
+            # Assume batch is a tuple of (inputs, targets)
+            inputs, targets = batch
+            
+        return inputs, targets
+    
+    def _analyze_model_weights(self):
+        """Analyze model weights for patterns and optimizations"""
+        # Save model to temporary safetensor file for analysis
+        model_path = "temp_model.safetensors"
+        tensors = {name: param for name, param in self.model.named_parameters()}
+        save_file(tensors, model_path)
+        
+        # Analyze weights
+        results = self.analyzer.analyze_safetensor_weights(model_path)
+        
+        # Log findings
+        logger.info(f"Weight Analysis: Found {len(results['shared_patterns'])} shared patterns")
+        logger.info(f"Potential memory savings: "
+                    f"{sum(i['potential_savings'] for i in results['memory_impact'].values())/1024/1024:.2f}MB")
+        
+        # Clean up
+        if os.path.exists(model_path):
+            os.remove(model_path)
+
+@dataclass
+class ModelConfig:
+    name: str
+    model_id: str
+    tokenizer_id: str
+    
+CONFIGS = {
+    "7b": ModelConfig(
+        name="7b",
+        model_id="scrapegoat/ScrapeGoat-Music-Stage1",
+        tokenizer_id="scrapegoat/ScrapeGoat-Music-Stage1"
+    ),
+    "1b": ModelConfig(
+        name="1b",
+        model_id="scrapegoat/ScrapeGoat-Music-Stage2",
+        tokenizer_id="scrapegoat/ScrapeGoat-Music-Stage2"
+    )
+}
+
+class MusicFineTuner:
+    def __init__(
+        self,
+        model_size: str,
+        dataset_path: str,
+        output_dir: str,
+        device: str = "auto",
+        batch_size: int = 4,
+        gradient_accumulation_steps: int = 4,
+        learning_rate: float = 1e-5,
+        num_epochs: int = 3,
+        use_hcf: bool = True
+    ):
+        self.config = CONFIGS[model_size]
+        self.dataset_path = Path(dataset_path)
+        self.output_dir = Path(output_dir)
+        self.device = self._setup_device(device)
+        self.use_hcf = use_hcf
+        self.training_args = TrainingArguments(
+            output_dir=str(self.output_dir),
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            learning_rate=learning_rate,
+            num_train_epochs=num_epochs,
+            logging_steps=100,
+            save_steps=1000,
+            evaluation_strategy="steps",
+            eval_steps=500,
+            save_total_limit=3,
+            load_best_model_at_end=True,
+            gradient_checkpointing=True,
+            fp16=torch.cuda.is_available(),
+            optim="adamw_torch"
+        )
+        
+    def _setup_device(self, device: str) -> str:
+        if device == "auto":
+            if torch.cuda.is_available():
+                return "cuda"
+            elif torch.backends.mps.is_available():
+                return "mps"
+            else:
+                return "cpu"
+        return device
+    
+    def _load_model_and_tokenizer(self):
+        logger.info(f"Loading model {self.config.model_id}")
+        
+        # Determine dtype based on device
+        dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            self.config.model_id,
+            torch_dtype=dtype,
+            device_map="auto" if self.device == "cuda" else None,
+            attn_implementation="flash_attention_2" if self.device == "cuda" else "eager"
+        )
+        
+        tokenizer = AutoTokenizer.from_pretrained(self.config.tokenizer_id)
+        return model, tokenizer
+    
+    def _prepare_dataset(self, tokenizer):
+        logger.info("Preparing dataset")
+        
+        with open(self.dataset_path / "metadata" / "dataset_info.json") as f:
+            metadata = json.load(f)
+        
+        def generate_text(item):
+            return f"Genre: {item['genre']}\nDuration: {item['duration']:.2f}s\nTitle: {item['title']}\nArtist: {item['artist']}\n"
+        
+        texts = [generate_text(item) for item in metadata["files"]]
+        dataset = Dataset.from_dict({"text": texts})
+        
+        def tokenize(examples):
+            return tokenizer(
+                examples["text"],
+                truncation=True,
+                padding="max_length",
+                max_length=512,
+                return_tensors="pt"
+            )
+        
+        tokenized_dataset = dataset.map(
+            tokenize,
+            batched=True,
+            remove_columns=dataset.column_names
+        )
+        
+        return tokenized_dataset
+    
+    def train(self):
+        # Create output directory
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Load model and tokenizer
+        model, tokenizer = self._load_model_and_tokenizer()
+        
+        # Prepare dataset
+        dataset = self._prepare_dataset(tokenizer)
+        
+        # Split dataset
+        dataset = dataset.train_test_split(test_size=0.1)
+        
+        if self.use_hcf:
+            logger.info("Using HCF-aware training")
+            # Create custom HCF optimizer
+            optimizer = HCFTrainingOptimizer(
+                model.parameters(),
+                lr=self.training_args.learning_rate,
+                weight_quantization=True,
+                maintain_patterns=True
+            )
+            
+            # Create HCF trainer
+            hcf_trainer = HCFAwareTrainer(model, optimizer)
+            
+            # Create custom training loop
+            train_loader = torch.utils.data.DataLoader(
+                dataset["train"],
+                batch_size=self.training_args.per_device_train_batch_size,
+                shuffle=True
+            )
+            
+            # Training loop with HCF awareness
+            criterion = torch.nn.CrossEntropyLoss()
+            for epoch in range(int(self.training_args.num_train_epochs)):
+                stats = hcf_trainer.train_epoch(train_loader, criterion, epoch)
+                
+                # Log training metrics
+                logger.info(f"Epoch {epoch} completed")
+                logger.info(f"Memory Savings: {stats.memory_savings/1024/1024:.2f}MB")
+                logger.info(f"Quantization Error: {stats.quantization_error:.6f}")
+                logger.info(f"Convergence Rate: {stats.convergence_rate:.4f}")
+                
+                # Save checkpoint
+                self._save_hcf_checkpoint(model, tokenizer, epoch)
+        else:
+            # Use standard HuggingFace Trainer
+            logger.info("Using standard training")
+            trainer = Trainer(
+                model=model,
+                args=self.training_args,
+                train_dataset=dataset["train"],
+                eval_dataset=dataset["test"],
+                data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
+            )
+            
+            # Train
+            logger.info("Starting training")
+            trainer.train()
+        
+        # Save final model
+        logger.info("Saving model")
+        model.save_pretrained(str(self.output_dir / "final_model"))
+        tokenizer.save_pretrained(str(self.output_dir / "final_model"))
+    
+    def _save_hcf_checkpoint(self, model, tokenizer, epoch):
+        """Save checkpoint with HCF metadata"""
+        checkpoint_dir = self.output_dir / f"checkpoint-{epoch}"
+        checkpoint_dir.mkdir(exist_ok=True)
+        
+        # Save model and tokenizer
+        model.save_pretrained(str(checkpoint_dir))
+        tokenizer.save_pretrained(str(checkpoint_dir))
+        
+        # Analyze and save HCF metadata
+        analyzer = SafeTensorHCFAnalyzer()
+        
+        # Save tensors to analyze
+        model_path = str(checkpoint_dir / "model.safetensors")
+        if os.path.exists(model_path):
+            results = analyzer.analyze_safetensor_weights(model_path)
+            
+            # Save analysis results
+            with open(checkpoint_dir / "hcf_analysis.json", "w") as f:
+                json.dump(results, f, indent=2)
+        
+        logger.info(f"Saved checkpoint at {checkpoint_dir}")
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_size", type=str, choices=["1b", "7b"], required=True)
+    parser.add_argument("--dataset_path", type=str, required=True)
+    parser.add_argument("--output_dir", type=str, required=True)
+    parser.add_argument("--device", type=str, default="auto")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=4)
+    parser.add_argument("--learning_rate", type=float, default=1e-5)
+    parser.add_argument("--num_epochs", type=int, default=3)
+    parser.add_argument("--use_hcf", action="store_true", help="Enable HCF-aware training")
+    args = parser.parse_args()
+    
+    fine_tuner = MusicFineTuner(
+        model_size=args.model_size,
+        dataset_path=args.dataset_path,
+        output_dir=args.output_dir,
+        device=args.device,
+        batch_size=args.batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        num_epochs=args.num_epochs,
+        use_hcf=args.use_hcf
+    )
+    fine_tuner.train() 

train_local.py

@@ -0,0 +1,331 @@
+#!/usr/bin/env python3
+"""
+Training script for ScrapeGoat Music models using local model files with HCF optimization.
+Optimized for local training with the models in the provided directory structure.
+"""
+
+import os
+import sys
+import json
+import torch
+import logging
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Tuple, Any
+import transformers
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from datasets import Dataset
+import numpy as np
+from accelerate import Accelerator
+from safetensors import safe_open
+from safetensors.torch import save_file, load_file
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, 
+                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Add xcodec_mini_infer to path to access its modules
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+XCODEC_PATH = os.path.join(SCRIPT_DIR, "xcodec_mini_infer")
+sys.path.append(XCODEC_PATH)
+
+# Import HCF training components from train_hcf.py
+from train_hcf import (
+    TensorInfo, 
+    SafeTensorHCFAnalyzer, 
+    TrainingStatistics,
+    HCFTrainingOptimizer, 
+    HCFAwareTrainer
+)
+
+@dataclass
+class LocalModelConfig:
+    """Configuration for local model directories"""
+    model_path: str
+    name: str
+    
+    @property
+    def model_dir(self) -> str:
+        return os.path.abspath(self.model_path)
+
+class LocalFineTuner:
+    """Fine-tuner that works with local model files"""
+    
+    def __init__(
+        self,
+        model_config: LocalModelConfig,
+        dataset_path: str,
+        output_dir: str,
+        device: str = "auto",
+        batch_size: int = 4,
+        gradient_accumulation_steps: int = 4,
+        learning_rate: float = 1e-5,
+        num_epochs: int = 3,
+        use_hcf: bool = True
+    ):
+        self.model_config = model_config
+        self.dataset_path = Path(dataset_path)
+        self.output_dir = Path(output_dir)
+        self.device = self._setup_device(device)
+        self.use_hcf = use_hcf
+        
+        # Ensure output directory exists
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Set up training arguments
+        self.training_args = TrainingArguments(
+            output_dir=str(self.output_dir),
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            learning_rate=learning_rate,
+            num_train_epochs=num_epochs,
+            logging_steps=100,
+            save_steps=1000,
+            evaluation_strategy="steps",
+            eval_steps=500,
+            save_total_limit=3,
+            load_best_model_at_end=True,
+            gradient_checkpointing=True,
+            fp16=torch.cuda.is_available(),
+            optim="adamw_torch"
+        )
+        
+    def _setup_device(self, device: str) -> str:
+        """Set up the training device"""
+        if device == "auto":
+            if torch.cuda.is_available():
+                return "cuda"
+            elif torch.backends.mps.is_available():
+                return "mps"
+            else:
+                return "cpu"
+        return device
+    
+    def _load_model_and_tokenizer(self):
+        """Load model and tokenizer from local path"""
+        logger.info(f"Loading model from {self.model_config.model_dir}")
+        
+        # Determine dtype based on device
+        dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
+        
+        # Load model from local path
+        model = AutoModelForCausalLM.from_pretrained(
+            self.model_config.model_dir,
+            torch_dtype=dtype,
+            device_map="auto" if self.device == "cuda" else None,
+            attn_implementation="flash_attention_2" if self.device == "cuda" else "eager",
+            local_files_only=True
+        )
+        
+        # Load tokenizer from local path
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.model_config.model_dir,
+            local_files_only=True
+        )
+        
+        return model, tokenizer
+    
+    def _prepare_dataset(self, tokenizer):
+        """Prepare dataset for training"""
+        logger.info("Preparing dataset")
+        
+        # Load metadata
+        with open(self.dataset_path / "metadata" / "dataset_info.json") as f:
+            metadata = json.load(f)
+        
+        # Define text generation from metadata
+        def generate_text(item):
+            return f"Genre: {item['genre']}\nDuration: {item['duration']:.2f}s\nTitle: {item['title']}\nArtist: {item['artist']}\n"
+        
+        # Generate text samples
+        texts = [generate_text(item) for item in metadata["files"]]
+        dataset = Dataset.from_dict({"text": texts})
+        
+        # Tokenize function
+        def tokenize(examples):
+            return tokenizer(
+                examples["text"],
+                truncation=True,
+                padding="max_length",
+                max_length=512,
+                return_tensors="pt"
+            )
+        
+        # Apply tokenization
+        tokenized_dataset = dataset.map(
+            tokenize,
+            batched=True,
+            remove_columns=dataset.column_names
+        )
+        
+        return tokenized_dataset
+    
+    def train(self):
+        """Train the model with HCF optimization"""
+        # Create output directory
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Log training configuration
+        logger.info(f"Training {self.model_config.name} model with HCF optimization")
+        logger.info(f"Model path: {self.model_config.model_dir}")
+        logger.info(f"Dataset path: {self.dataset_path}")
+        logger.info(f"Output directory: {self.output_dir}")
+        logger.info(f"Device: {self.device}")
+        logger.info(f"HCF optimization: {'enabled' if self.use_hcf else 'disabled'}")
+        
+        # Load model and tokenizer
+        model, tokenizer = self._load_model_and_tokenizer()
+        
+        # Prepare dataset
+        dataset = self._prepare_dataset(tokenizer)
+        
+        # Split dataset
+        dataset = dataset.train_test_split(test_size=0.1)
+        
+        if self.use_hcf:
+            logger.info("Using HCF-aware training")
+            # Create custom HCF optimizer
+            optimizer = HCFTrainingOptimizer(
+                model.parameters(),
+                lr=self.training_args.learning_rate,
+                weight_quantization=True,
+                maintain_patterns=True
+            )
+            
+            # Create HCF trainer
+            hcf_trainer = HCFAwareTrainer(model, optimizer)
+            
+            # Create custom training loop
+            train_loader = torch.utils.data.DataLoader(
+                dataset["train"],
+                batch_size=self.training_args.per_device_train_batch_size,
+                shuffle=True
+            )
+            
+            # Training loop with HCF awareness
+            criterion = torch.nn.CrossEntropyLoss()
+            for epoch in range(int(self.training_args.num_train_epochs)):
+                stats = hcf_trainer.train_epoch(train_loader, criterion, epoch)
+                
+                # Log training metrics
+                logger.info(f"Epoch {epoch} completed")
+                logger.info(f"Memory Savings: {stats.memory_savings/1024/1024:.2f}MB")
+                logger.info(f"Quantization Error: {stats.quantization_error:.6f}")
+                logger.info(f"Convergence Rate: {stats.convergence_rate:.4f}")
+                
+                # Save checkpoint
+                self._save_hcf_checkpoint(model, tokenizer, epoch)
+        else:
+            # Use standard HuggingFace Trainer
+            logger.info("Using standard training")
+            trainer = Trainer(
+                model=model,
+                args=self.training_args,
+                train_dataset=dataset["train"],
+                eval_dataset=dataset["test"],
+                data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
+            )
+            
+            # Train
+            logger.info("Starting training")
+            trainer.train()
+        
+        # Save final model
+        logger.info("Saving model")
+        final_output_dir = self.output_dir / "final_model"
+        final_output_dir.mkdir(exist_ok=True)
+        
+        model.save_pretrained(str(final_output_dir))
+        tokenizer.save_pretrained(str(final_output_dir))
+        
+        logger.info(f"Training complete. Model saved to {final_output_dir}")
+    
+    def _save_hcf_checkpoint(self, model, tokenizer, epoch):
+        """Save checkpoint with HCF metadata"""
+        checkpoint_dir = self.output_dir / f"checkpoint-{epoch}"
+        checkpoint_dir.mkdir(exist_ok=True)
+        
+        # Save model and tokenizer
+        model.save_pretrained(str(checkpoint_dir))
+        tokenizer.save_pretrained(str(checkpoint_dir))
+        
+        # Analyze and save HCF metadata
+        analyzer = SafeTensorHCFAnalyzer()
+        
+        # Save tensors to analyze
+        model_path = str(checkpoint_dir / "model.safetensors")
+        if os.path.exists(model_path):
+            results = analyzer.analyze_safetensor_weights(model_path)
+            
+            # Save analysis results
+            with open(checkpoint_dir / "hcf_analysis.json", "w") as f:
+                json.dump(results, f, indent=2)
+        
+        logger.info(f"Saved checkpoint at {checkpoint_dir}")
+
+def main():
+    """Main function for training"""
+    import argparse
+    parser = argparse.ArgumentParser(description="Retrain ScrapeGoat Music models with HCF optimization")
+    parser.add_argument("--model", type=str, choices=["7b", "1b"], required=True, 
+                       help="Model size to train")
+    parser.add_argument("--dataset_path", type=str, required=True,
+                       help="Path to processed dataset")
+    parser.add_argument("--output_dir", type=str, required=True,
+                       help="Directory to save trained model")
+    parser.add_argument("--device", type=str, default="auto",
+                       help="Device to use (cuda, cpu, mps, or auto)")
+    parser.add_argument("--batch_size", type=int, default=4,
+                       help="Batch size for training")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=4,
+                       help="Gradient accumulation steps")
+    parser.add_argument("--learning_rate", type=float, default=1e-5,
+                       help="Learning rate")
+    parser.add_argument("--num_epochs", type=int, default=3,
+                       help="Number of training epochs")
+    parser.add_argument("--use_hcf", action="store_true", default=True,
+                       help="Enable HCF optimization")
+    parser.add_argument("--base_dir", type=str, default=os.getcwd(),
+                       help="Base directory containing model folders")
+    
+    args = parser.parse_args()
+    
+    # Set up model configuration based on size
+    if args.model == "7b":
+        model_path = os.path.join(args.base_dir, "scrapegoat/ScrapeGoat-Music-Stage1")
+        model_config = LocalModelConfig(
+            model_path=model_path,
+            name="ScrapeGoatMusic 7B"
+        )
+    else:
+        model_path = os.path.join(args.base_dir, "scrapegoat/ScrapeGoat-Music-Stage2")
+        model_config = LocalModelConfig(
+            model_path=model_path,
+            name="ScrapeGoatMusic 1B"
+        )
+    
+    # Create fine-tuner
+    fine_tuner = LocalFineTuner(
+        model_config=model_config,
+        dataset_path=args.dataset_path,
+        output_dir=args.output_dir,
+        device=args.device,
+        batch_size=args.batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        num_epochs=args.num_epochs,
+        use_hcf=args.use_hcf
+    )
+    
+    # Train model
+    fine_tuner.train()
+
+if __name__ == "__main__":
+    main()