Create tokenizer.py

to be developed further

tokenizer.py

@@ -0,0 +1,512 @@
+"""
+Tokenization for VCF data with support for hierarchical structures
+"""
+
+import json
+import pickle
+import logging
+from pathlib import Path
+from collections import defaultdict, Counter
+from typing import Dict, List, Tuple, Optional, Union, Any
+import numpy as np
+
+from transformers import PreTrainedTokenizer
+from transformers.tokenization_utils import AddedToken
+
+from config import DataConfig, ConfigManager
+from parser import MutationRecord
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class HierarchicalVCFTokenizer(PreTrainedTokenizer):
+    
+    vocab_files_names = {
+        "vocab_file": "vocab.json",
+        "mutation_vocab_file": "mutation_vocab.json"
+    }
+    
+    def __init__(self,
+                 vocab_file: Optional[str] = None,
+                 mutation_vocab_file: Optional[str] = None,
+                 config: Optional[DataConfig] = None,
+                 **kwargs):
+        
+        # Initialize special tokens
+        self.config = config or DataConfig()
+        
+        # Set up special tokens
+        special_tokens = self.config.special_tokens
+        pad_token = special_tokens.get("pad_token", "[PAD]")
+        unk_token = special_tokens.get("unk_token", "[UNK]")
+        sep_token = special_tokens.get("sep_token", "[SEP]")
+        cls_token = special_tokens.get("cls_token", "[CLS]")
+        
+        super().__init__(
+            pad_token=pad_token,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            cls_token=cls_token,
+            **kwargs
+        )
+        
+        # Initialize vocabularies for different mutation fields
+        self.mutation_fields = ['impact', 'ref', 'alt', 'chromosome', 'pathway', 'gene']
+        self.field_vocabs = {}
+        
+        # Initialize vocabularies
+        self._initialize_vocabularies()
+        
+        # Load existing vocabularies if provided
+        if vocab_file and Path(vocab_file).exists():
+            self.load_vocabulary(vocab_file)
+        
+        if mutation_vocab_file and Path(mutation_vocab_file).exists():
+            self.load_mutation_vocabulary(mutation_vocab_file)
+        
+        # Statistics
+        self.tokenization_stats = {
+            'total_samples': 0,
+            'total_mutations': 0,
+            'vocab_sizes': {}
+        }
+    
+    def _initialize_vocabularies(self) -> None:
+        for field in self.mutation_fields:
+            self.field_vocabs[field] = {
+                self.pad_token: 0,
+                self.unk_token: 1,
+                self.sep_token: 2,
+                self.cls_token: 3
+            }
+        
+        # Add common genomic tokens
+        self._add_common_genomic_tokens()
+    
+    def _add_common_genomic_tokens(self) -> None:
+        """To be made scalable and dynamic"""
+        # Common impact values
+        common_impacts = ["HIGH", "MODERATE", "LOW", "MODIFIER"]
+        for impact in common_impacts:
+            if impact not in self.field_vocabs['impact']:
+                self.field_vocabs['impact'][impact] = len(self.field_vocabs['impact'])
+        
+        # Common nucleotides
+        nucleotides = ["A", "T", "G", "C", "N", "-"]
+        for nt in nucleotides:
+            for field in ['ref', 'alt']:
+                if nt not in self.field_vocabs[field]:
+                    self.field_vocabs[field][nt] = len(self.field_vocabs[field])
+        
+        # Common chromosomes
+        chromosomes = [str(i) for i in range(1, 23)] + ["X", "Y", "MT"]
+        for chrom in chromosomes:
+            if chrom not in self.field_vocabs['chromosome']:
+                self.field_vocabs['chromosome'][chrom] = len(self.field_vocabs['chromosome'])
+    
+    def build_vocabulary(self, hierarchical_data: Dict[str, Any]) -> None:
+        """
+        Args:
+            hierarchical_data: Parsed VCF data structure
+        """
+        logger.info("Building vocabularies from hierarchical data...")
+        
+        vocab_counters = {field: Counter() for field in self.mutation_fields}
+        
+        for sample_id, pathways in hierarchical_data.items():
+            for pathway_id, chromosomes in pathways.items():
+                # Count pathway occurrences
+                vocab_counters['pathway'][pathway_id] += 1
+                
+                for chrom_id, genes in chromosomes.items():
+                    # Count chromosome occurrences
+                    vocab_counters['chromosome'][chrom_id] += 1
+                    
+                    for gene_id, mutations in genes.items():
+                        # Count gene occurrences
+                        vocab_counters['gene'][gene_id] += 1
+                        
+                        for mutation in mutations:
+                            if isinstance(mutation, MutationRecord):
+                                # Count mutation field values
+                                vocab_counters['impact'][mutation.impact] += 1
+                                vocab_counters['ref'][mutation.reference] += 1
+                                vocab_counters['alt'][mutation.alternate] += 1
+                            elif isinstance(mutation, dict):
+                                # Handle dictionary format
+                                vocab_counters['impact'][mutation.get('impact', self.unk_token)] += 1
+                                vocab_counters['ref'][mutation.get('reference', self.unk_token)] += 1
+                                vocab_counters['alt'][mutation.get('alternate', self.unk_token)] += 1
+        
+        # Build vocabularies from counters
+        for field, counter in vocab_counters.items():
+            for token, count in counter.most_common():
+                if token and token not in self.field_vocabs[field]:
+                    self.field_vocabs[field][token] = len(self.field_vocabs[field])
+        
+        # Update statistics
+        self.tokenization_stats['vocab_sizes'] = {
+            field: len(vocab) for field, vocab in self.field_vocabs.items()
+        }
+        
+        logger.info(f"Vocabulary sizes: {self.tokenization_stats['vocab_sizes']}")
+    
+    def encode_hierarchical_sample(self, sample_data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Encode a single hierarchical sample into tokenized format.  
+        Args:
+            sample_data: Single sample from hierarchical data           
+        Returns:
+            Encoded sample with tokenized values
+        """
+        encoded_sample = {}
+        
+        for pathway_id, chromosomes in sample_data.items():
+            # Tokenize pathway ID
+            pathway_token = self.field_vocabs['pathway'].get(
+                pathway_id, self.field_vocabs['pathway'][self.unk_token]
+            )
+            
+            encoded_sample[pathway_token] = {}
+            
+            for chrom_id, genes in chromosomes.items():
+                # Tokenize chromosome ID
+                chrom_token = self.field_vocabs['chromosome'].get(
+                    chrom_id, self.field_vocabs['chromosome'][self.unk_token]
+                )
+                
+                encoded_sample[pathway_token][chrom_token] = {}
+                
+                for gene_id, mutations in genes.items():
+                    # Tokenize gene ID
+                    gene_token = self.field_vocabs['gene'].get(
+                        gene_id, self.field_vocabs['gene'][self.unk_token]
+                    )
+                    
+                    # Encode mutations
+                    encoded_mutations = self._encode_mutations(mutations)
+                    encoded_sample[pathway_token][chrom_token][gene_token] = encoded_mutations
+        
+        return encoded_sample
+    
+    def _encode_mutations(self, mutations: List[Union[MutationRecord, Dict]]) -> Dict[str, List[int]]:
+        encoded_mutations = {
+            'impact': [],
+            'ref': [],
+            'alt': []
+        }
+        
+        for mutation in mutations:
+            if isinstance(mutation, MutationRecord):
+                impact = mutation.impact
+                ref = mutation.reference
+                alt = mutation.alternate
+            elif isinstance(mutation, dict):
+                impact = mutation.get('impact', self.unk_token)
+                ref = mutation.get('reference', self.unk_token)
+                alt = mutation.get('alternate', self.unk_token)
+            else:
+                continue
+            
+            # Tokenize each field
+            encoded_mutations['impact'].append(
+                self.field_vocabs['impact'].get(impact, self.field_vocabs['impact'][self.unk_token])
+            )
+            encoded_mutations['ref'].append(
+                self.field_vocabs['ref'].get(ref, self.field_vocabs['ref'][self.unk_token])
+            )
+            encoded_mutations['alt'].append(
+                self.field_vocabs['alt'].get(alt, self.field_vocabs['alt'][self.unk_token])
+            )
+        
+        return encoded_mutations
+    
+    def encode_batch(self, batch_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Encode a batch of hierarchical samples.        
+        Args:
+            batch_data: List of sample dictionaries            
+        Returns:
+            List of encoded samples
+        """
+        encoded_batch = []
+        
+        for sample_data in batch_data:
+            encoded_sample = self.encode_hierarchical_sample(sample_data)
+            encoded_batch.append(encoded_sample)
+        
+        self.tokenization_stats['total_samples'] += len(batch_data)
+        
+        return encoded_batch
+    
+    def decode_tokens(self, field: str, token_ids: List[int]) -> List[str]:
+        """
+        Decode token IDs back to original values.       
+        Args:
+            field: Field name ('impact', 'ref', 'alt', etc.)
+            token_ids: List of token IDs           
+        Returns:
+            List of decoded tokens
+        """
+        if field not in self.field_vocabs:
+            raise ValueError(f"Unknown field: {field}")
+        
+        id_to_token = {v: k for k, v in self.field_vocabs[field].items()}
+        return [id_to_token.get(token_id, self.unk_token) for token_id in token_ids]
+    
+    def get_vocab_size(self, field: str) -> int:
+        """Get vocabulary size for a specific field."""
+        if field not in self.field_vocabs:
+            raise ValueError(f"Unknown field: {field}")
+        return len(self.field_vocabs[field])
+    
+    def get_all_vocab_sizes(self) -> Dict[str, int]:
+        """Get vocabulary sizes for all fields."""
+        return {field: len(vocab) for field, vocab in self.field_vocabs.items()}
+    
+    def save_vocabulary(self, save_directory: Union[str, Path], filename_prefix: Optional[str] = None) -> Tuple[str, ...]:
+        """
+        Args:
+            save_directory: Directory to save vocabularies
+            filename_prefix: Optional prefix for filenames
+            
+        Returns:
+            Tuple of saved file paths
+        """
+        save_directory = Path(save_directory)
+        save_directory.mkdir(parents=True, exist_ok=True)
+        
+        prefix = f"{filename_prefix}_" if filename_prefix else ""
+        
+        # Save mutation vocabularies
+        mutation_vocab_file = save_directory / f"{prefix}mutation_vocab.json"
+        with open(mutation_vocab_file, 'w') as f:
+            json.dump(self.field_vocabs, f, indent=2)
+        
+        # Save tokenizer configuration
+        config_file = save_directory / f"{prefix}tokenizer_config.json"
+        config_data = {
+            'tokenizer_class': self.__class__.__name__,
+            'special_tokens': {
+                'pad_token': self.pad_token,
+                'unk_token': self.unk_token,
+                'sep_token': self.sep_token,
+                'cls_token': self.cls_token
+            },
+            'vocab_sizes': self.get_all_vocab_sizes(),
+            'mutation_fields': self.mutation_fields
+        }
+        
+        with open(config_file, 'w') as f:
+            json.dump(config_data, f, indent=2)
+        
+        logger.info(f"Vocabularies saved to {save_directory}")
+        
+        return str(mutation_vocab_file), str(config_file)
+    
+    def load_vocabulary(self, vocab_file: Union[str, Path]) -> None:
+        vocab_file = Path(vocab_file)
+        
+        if not vocab_file.exists():
+            raise FileNotFoundError(f"Vocabulary file not found: {vocab_file}")
+        
+        with open(vocab_file, 'r') as f:
+            vocab_data = json.load(f)
+        
+        # Update vocabularies
+        for field, vocab in vocab_data.items():
+            if field in self.mutation_fields:
+                self.field_vocabs[field] = vocab
+        
+        logger.info(f"Vocabularies loaded from {vocab_file}")
+    
+    def load_mutation_vocabulary(self, mutation_vocab_file: Union[str, Path]) -> None:
+        """Load mutation-specific vocabularies from file."""
+        self.load_vocabulary(mutation_vocab_file)
+    
+    def create_padding_masks(self, encoded_sample: Dict[str, Any], max_lengths: Dict[str, int]) -> Dict[str, Any]:
+        """
+        Create padding masks for hierarchical data.      
+        Args:
+            encoded_sample: Encoded sample data
+            max_lengths: Maximum lengths for each level   
+        Returns:
+            Sample with padding masks
+        """
+        masked_sample = {}
+        
+        for pathway_token, chromosomes in encoded_sample.items():
+            masked_sample[pathway_token] = {}
+            
+            for chrom_token, genes in chromosomes.items():
+                masked_sample[pathway_token][chrom_token] = {}
+                
+                for gene_token, mutations in genes.items():
+                    masked_mutations = {}
+                    
+                    for field, token_list in mutations.items():
+                        max_len = max_lengths.get(f'mutations_{field}', 100)
+                        
+                        # Pad or truncate
+                        if len(token_list) < max_len:
+                            padded_list = token_list + [self.field_vocabs[field][self.pad_token]] * (max_len - len(token_list))
+                            mask = [1] * len(token_list) + [0] * (max_len - len(token_list))
+                        else:
+                            padded_list = token_list[:max_len]
+                            mask = [1] * max_len
+                        
+                        masked_mutations[field] = {
+                            'tokens': padded_list,
+                            'mask': mask
+                        }
+                    
+                    masked_sample[pathway_token][chrom_token][gene_token] = masked_mutations
+        
+        return masked_sample
+    
+    def get_tokenization_statistics(self) -> Dict[str, Any]:
+        stats = self.tokenization_stats.copy()
+        stats['vocab_sizes'] = self.get_all_vocab_sizes()
+        return stats
+    
+    # Hugging Face compatibility methods
+    @property
+    def vocab_size(self) -> int:
+        return sum(len(vocab) for vocab in self.field_vocabs.values())
+    
+    def get_vocab(self) -> Dict[str, int]:
+        combined_vocab = {}
+        offset = 0
+        
+        for field, vocab in self.field_vocabs.items():
+            for token, idx in vocab.items():
+                combined_vocab[f"{field}:{token}"] = idx + offset
+            offset += len(vocab)
+        
+        return combined_vocab
+    
+    def _tokenize(self, text: str) -> List[str]:
+        # This is a simplified implementation for compatibility
+        # In practice, hierarchical data should be processed differently
+        return text.split()
+    
+    def _convert_token_to_id(self, token: str) -> int:
+        # Parse field:token format
+        if ':' in token:
+            field, actual_token = token.split(':', 1)
+            if field in self.field_vocabs:
+                return self.field_vocabs[field].get(actual_token, self.field_vocabs[field][self.unk_token])
+        
+        return self.field_vocabs.get('impact', {}).get(self.unk_token, 1)
+    
+    def _convert_id_to_token(self, index: int) -> str:
+        # This is a simplified reverse lookup
+        for field, vocab in self.field_vocabs.items():
+            id_to_token = {v: k for k, v in vocab.items()}
+            if index in id_to_token:
+                return f"{field}:{id_to_token[index]}"
+        
+        return self.unk_token
+
+
+class HierarchicalDataCollator:
+  
+    def __init__(self, tokenizer: HierarchicalVCFTokenizer, max_lengths: Optional[Dict[str, int]] = None):
+        self.tokenizer = tokenizer
+        self.max_lengths = max_lengths or {
+            'mutations_impact': 50,
+            'mutations_ref': 50,
+            'mutations_alt': 50,
+            'genes_per_chromosome': 100,
+            'chromosomes_per_pathway': 25,
+            'pathways_per_sample': 50
+        }
+    
+    def __call__(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Collate batch of hierarchical samples.
+        Args:
+            batch: List of encoded hierarchical samples
+        Returns:
+            Collated batch ready for model input
+        """
+        collated_batch = {
+            'samples': [],
+            'batch_size': len(batch),
+            'metadata': {
+                'num_pathways': [],
+                'num_chromosomes': [],
+                'num_genes': [],
+                'num_mutations': []
+            }
+        }
+        
+        for sample in batch:
+            # Create padding masks
+            masked_sample = self.tokenizer.create_padding_masks(sample, self.max_lengths)
+            collated_batch['samples'].append(masked_sample)
+            
+            # Collect metadata
+            num_pathways = len(sample)
+            num_chromosomes = sum(len(chroms) for chroms in sample.values())
+            num_genes = sum(
+                len(genes) for chroms in sample.values() 
+                for genes in chroms.values()
+            )
+            num_mutations = sum(
+                len(mutations.get('impact', [])) 
+                for chroms in sample.values() 
+                for genes in chroms.values() 
+                for mutations in genes.values()
+            )
+            
+            collated_batch['metadata']['num_pathways'].append(num_pathways)
+            collated_batch['metadata']['num_chromosomes'].append(num_chromosomes)
+            collated_batch['metadata']['num_genes'].append(num_genes)
+            collated_batch['metadata']['num_mutations'].append(num_mutations)
+        
+        return collated_batch
+
+
+def create_tokenizer_from_config(config_manager: ConfigManager) -> HierarchicalVCFTokenizer:
+    """Create tokenizer from configuration manager."""
+    return HierarchicalVCFTokenizer(config=config_manager.data_config)
+
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Example usage
+    config_manager = ConfigManager()
+    tokenizer = create_tokenizer_from_config(config_manager)
+    
+    # Example hierarchical data structure
+    example_data = {
+        'sample1': {
+            'pathway1': {
+                'chr1': {
+                    'gene1': [
+                        {
+                            'impact': 'HIGH',
+                            'reference': 'A',
+                            'alternate': 'T'
+                        }
+                    ]
+                }
+            }
+        }
+    }
+    
+    # Build vocabulary
+    tokenizer.build_vocabulary({'sample1': example_data['sample1']})
+    
+    # Encode sample
+    encoded = tokenizer.encode_hierarchical_sample(example_data['sample1'])
+    print(f"Encoded sample: {encoded}")
+    
+    # Save vocabulary
+    tokenizer.save_vocabulary("./tokenizer_files")
+    
+    print(f"Tokenization statistics: {tokenizer.get_tokenization_statistics()}")