Create custom_tokenizer.py

custom_tokenizer.py

@@ -0,0 +1,437 @@
+import os
+import json
+import pickle
+import argparse
+from collections import Counter, defaultdict
+from typing import List, Dict, Set, Optional, Tuple
+import re
+import unicodedata
+class TechnicalTokenizer:
+    """
+    Custom tokenizer optimized for technical content and conversations
+    """    
+    def __init__(self, vocab_size: int = 32000, min_freq: int = 2):
+        self.vocab_size = vocab_size
+        self.min_freq = min_freq       
+        self.special_tokens = {
+            '<pad>': 0,
+            '<unk>': 1,
+            '<bos>': 2,
+            '<eos>': 3,
+            '<system>': 4,
+            '<user>': 5,
+            '<assistant>': 6,
+            '<|endoftext|>': 7,
+            '<|newline|>': 8,
+            '<|tab|>': 9,
+            '<|code|>': 10,
+            '<|/code|>': 11,
+            '<|math|>': 12,
+            '<|/math|>': 13
+        }       
+        self.vocab = {}
+        self.id_to_token = {}
+        self.token_frequencies = Counter()
+        self.bpe_merges = []
+        self.bpe_cache = {}       
+        self.code_pattern = re.compile(r'```[\s\S]*?```|`[^`]+`')
+        self.url_pattern = re.compile(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+')
+        self.email_pattern = re.compile(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b')
+        self.number_pattern = re.compile(r'\b\d+\.?\d*\b')        
+        self.technical_terms = {
+            'function', 'variable', 'array', 'object', 'class', 'method', 'parameter',
+            'return', 'import', 'export', 'async', 'await', 'promise', 'callback',
+            'algorithm', 'datatype', 'boolean', 'integer', 'string', 'float',
+            'javascript', 'python', 'java', 'cpp', 'html', 'css', 'sql',
+            'api', 'json', 'xml', 'http', 'https', 'rest', 'graphql',
+            'equation', 'formula', 'theorem', 'proof', 'hypothesis',
+            'derivative', 'integral', 'matrix', 'vector', 'polynomial',
+            'probability', 'statistics', 'correlation', 'regression',
+            'neural', 'network', 'model', 'training', 'validation', 'test',
+            'accuracy', 'precision', 'recall', 'f1score', 'loss', 'gradient',
+            'backpropagation', 'forward', 'layer', 'neuron', 'weight', 'bias',
+            'transformer', 'attention', 'embedding', 'tokenization',
+            'database', 'server', 'client', 'protocol', 'encryption', 'security',
+            'authentication', 'authorization', 'deployment', 'docker', 'kubernetes',
+            'microservice', 'architecture', 'scalability', 'performance'
+        }       
+        self._init_vocab()   
+    def _init_vocab(self):
+        self.vocab = self.special_tokens.copy()
+        self.id_to_token = {v: k for k, v in self.special_tokens.items()}    
+    def normalize_text(self, text: str) -> str:
+        text = re.sub(r'\r\n|\r', '\n', text)
+        text = re.sub(r'\t', '<|tab|>', text)
+        text = unicodedata.normalize('NFKC', text)       
+        code_blocks = []
+        def replace_code(match):
+            code_blocks.append(match.group())
+            return f'<|code|>CODE_BLOCK_{len(code_blocks)-1}<|/code|>'       
+        text = self.code_pattern.sub(replace_code, text)
+        text = self.url_pattern.sub('<URL>', text)
+        text = self.email_pattern.sub('<EMAIL>', text)       
+        for i, code_block in enumerate(code_blocks):
+            text = text.replace(f'<|code|>CODE_BLOCK_{i}<|/code|>', code_block)        
+        return text    
+    def pre_tokenize(self, text: str) -> List[str]:
+        text = self.normalize_text(text)        
+        text = re.sub(r'<\|system\|>', ' <system> ', text)
+        text = re.sub(r'<\|user\|>', ' <user> ', text)
+        text = re.sub(r'<\|assistant\|>', ' <assistant> ', text)
+        text = re.sub(r'<\|endoftext\|>', ' <|endoftext|> ', text)       
+        tokens = re.findall(r'''
+            <[^>]+>|                    # Special tokens
+            \b\w+@\w+\.\w+\b|          # Email-like patterns  
+            https?://\S+|              # URLs
+            ```[\s\S]*?```|            # Code blocks
+            `[^`]+`|                   # Inline code
+            \b\d+\.?\d*\b|             # Numbers
+            \b[a-zA-Z]+(?:'[a-z]*)?|   # Words with optional apostrophes
+            [^\w\s]                    # Punctuation
+        ''', text, re.VERBOSE)       
+        return [token.strip() for token in tokens if token.strip()]    
+    def get_pairs(self, word_freqs: Dict[Tuple[str, ...], int]) -> Counter:
+        pairs = Counter()
+        for word, freq in word_freqs.items():
+            if len(word) < 2:
+                continue
+            for i in range(len(word) - 1):
+                pair = (word[i], word[i + 1])
+                pairs[pair] += freq
+        return pairs   
+    def merge_symbols(self, pair: Tuple[str, str], word_freqs: Dict[Tuple[str, ...], int]) -> Dict[Tuple[str, ...], int]:
+        new_word_freqs = {}
+        bigram = pair
+        for word, freq in word_freqs.items():
+            new_word = []
+            i = 0
+            while i < len(word):
+                if i < len(word) - 1 and (word[i], word[i + 1]) == bigram:
+                    new_word.append(word[i] + word[i + 1])
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word_freqs[tuple(new_word)] = freq
+        return new_word_freqs   
+    def train_bpe(self, texts: List[str]) -> None:
+        print("Training BPE tokenizer...")
+        word_freqs = Counter()       
+        for i, text in enumerate(texts):
+            if i % 10000 == 0:
+                print(f"Processing text {i}/{len(texts)}")           
+            tokens = self.pre_tokenize(text)
+            for token in tokens:
+                char_seq = tuple(token)
+                if len(char_seq) > 0:
+                    word_freqs[char_seq] += 1        
+        print(f"Found {len(word_freqs)} unique word patterns")
+        word_freqs = {word: freq for word, freq in word_freqs.items() if freq >= self.min_freq}        
+        for term in self.technical_terms:
+            if (term,) in word_freqs:
+                word_freqs[(term,)] *= 10        
+        all_chars = set()
+        for word in word_freqs:
+            all_chars.update(word)        
+        for char in sorted(all_chars):
+            if char not in self.vocab:
+                self.vocab[char] = len(self.vocab)
+                self.id_to_token[len(self.id_to_token)] = char        
+        target_vocab_size = self.vocab_size - len(self.special_tokens)
+        num_merges = target_vocab_size - len(self.vocab)       
+        for i in range(num_merges):
+            if i % 1000 == 0:
+                print(f"BPE merge {i}/{num_merges}")            
+            pairs = self.get_pairs(word_freqs)
+            if not pairs:
+                break            
+            best_pair = pairs.most_common(1)[0][0]
+            word_freqs = self.merge_symbols(best_pair, word_freqs)
+            merged_token = best_pair[0] + best_pair[1]
+            if merged_token not in self.vocab:
+                self.vocab[merged_token] = len(self.vocab)
+                self.id_to_token[len(self.id_to_token)] = merged_token
+            self.bpe_merges.append(best_pair)       
+        print(f"BPE training complete. Final vocabulary size: {len(self.vocab)}")
+        for word, freq in word_freqs.items():
+            for token in word:
+                self.token_frequencies[token] += freq  
+    def apply_bpe(self, word: str) -> List[str]:
+        if word in self.bpe_cache:
+            return self.bpe_cache[word]        
+        tokens = list(word)
+        for merge in self.bpe_merges:
+            i = 0
+            while i < len(tokens) - 1:
+                if tokens[i] == merge[0] and tokens[i + 1] == merge[1]:
+                    tokens = tokens[:i] + [merge[0] + merge[1]] + tokens[i + 2:]
+                else:
+                    i += 1
+        self.bpe_cache[word] = tokens
+        return tokens   
+    def tokenize(self, text: str) -> List[str]:
+        pre_tokens = self.pre_tokenize(text)
+        final_tokens = []
+        for token in pre_tokens:
+            if token in self.special_tokens or token in self.vocab:
+                final_tokens.append(token)
+            else:
+                bpe_tokens = self.apply_bpe(token)
+                final_tokens.extend(bpe_tokens)
+        return final_tokens   
+    def encode_ids(self, text: str, add_special_tokens: bool = True) -> List[int]:
+        tokens = self.tokenize(text)
+        if add_special_tokens:
+            tokens = ['<bos>'] + tokens + ['<eos>']       
+        ids = []
+        for token in tokens:
+            ids.append(self.vocab.get(token, self.vocab['<unk>']))        
+        return ids   
+    def decode_ids(self, ids: List[int], skip_special_tokens: bool = True) -> str:
+        tokens = []
+        for id in ids:
+            token = self.id_to_token.get(id, '<unk>')
+            if skip_special_tokens and token in self.special_tokens:
+                continue
+            tokens.append(token)        
+        text = ''.join(tokens)
+        text = text.replace('<|tab|>', '\t')
+        text = text.replace('<|newline|>', '\n')
+        return text   
+    def save(self, save_dir: str):
+        os.makedirs(save_dir, exist_ok=True)
+        with open(os.path.join(save_dir, 'vocab.json'), 'w', encoding='utf-8') as f:
+            json.dump(self.vocab, f, indent=2, ensure_ascii=False)
+        with open(os.path.join(save_dir, 'merges.txt'), 'w', encoding='utf-8') as f:
+            for merge in self.bpe_merges:
+                f.write(f"{merge[0]} {merge[1]}\n")
+        config = {
+            'vocab_size': self.vocab_size,
+            'min_freq': self.min_freq,
+            'special_tokens': self.special_tokens,
+            'technical_terms': list(self.technical_terms)
+        }
+        with open(os.path.join(save_dir, 'tokenizer_config.json'), 'w', encoding='utf-8') as f:
+            json.dump(config, f, indent=2, ensure_ascii=False)
+        with open(os.path.join(save_dir, 'token_frequencies.pkl'), 'wb') as f:
+            pickle.dump(dict(self.token_frequencies), f)
+        print(f"Tokenizer saved to {save_dir}")    
+    def load(self, save_dir: str):
+        with open(os.path.join(save_dir, 'vocab.json'), 'r', encoding='utf-8') as f:
+            self.vocab = json.load(f)
+        self.id_to_token = {v: k for k, v in self.vocab.items()}
+        with open(os.path.join(save_dir, 'merges.txt'), 'r', encoding='utf-8') as f:
+            self.bpe_merges = [tuple(line.strip().split()) for line in f if line.strip()]
+        config_file = os.path.join(save_dir, 'tokenizer_config.json')
+        if os.path.exists(config_file):
+            with open(config_file, 'r', encoding='utf-8') as f:
+                config = json.load(f)
+                self.vocab_size = config.get('vocab_size', self.vocab_size)
+                self.min_freq = config.get('min_freq', self.min_freq)
+                if 'technical_terms' in config:
+                    self.technical_terms = set(config['technical_terms'])
+        freq_file = os.path.join(save_dir, 'token_frequencies.pkl')
+        if os.path.exists(freq_file):
+            with open(freq_file, 'rb') as f:
+                self.token_frequencies = Counter(pickle.load(f))
+        self.bpe_cache = {}
+        print(f"Tokenizer loaded from {save_dir}")
+        print(f"Vocabulary size: {len(self.vocab)}")
+        print(f"Number of BPE merges: {len(self.bpe_merges)}")    
+    def get_vocab_size(self) -> int:
+        return len(self.vocab)   
+    def get_token_frequency(self, token: str) -> int:
+        return self.token_frequencies.get(token, 0)   
+    def analyze_tokenization(self, text: str):
+        tokens = self.tokenize(text)
+        ids = self.encode_ids(text, add_special_tokens=False)
+        print(f"Original text: {text}")
+        print(f"Tokens: {tokens}")
+        print(f"Token IDs: {ids}")
+        print(f"Number of tokens: {len(tokens)}")
+        print(f"Compression ratio: {len(text.split())/len(tokens):.2f}")
+        return tokens, ids
+class ConversationDataset:
+    """Dataset class for handling conversation data with the custom tokenizer"""  
+    def __init__(self, data_file: str, tokenizer: TechnicalTokenizer, max_length: int = 512):
+        self.data_file = data_file
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.conversations = []
+        self.load_conversations()   
+    def load_conversations(self):
+        print(f"Loading conversations from {self.data_file}")       
+        if self.data_file.endswith('.jsonl'):
+            self.load_jsonl()
+        else:
+            self.load_text()        
+        print(f"Loaded {len(self.conversations)} conversations")   
+    def load_jsonl(self):
+        with open(self.data_file, 'r', encoding='utf-8') as f:
+            for line in f:
+                try:
+                    conv = json.loads(line.strip())
+                    messages = conv.get("messages", [])
+                    if not messages:
+                        continue                   
+                    text_parts = []
+                    for msg in messages:
+                        role = msg.get("role", "")
+                        content = msg.get("content", "").strip()
+                        if not content:
+                            continue
+                        if role == "system":
+                            continue
+                        elif role == "user":
+                            text_parts.append(f"<user> {content}")
+                        elif role == "assistant":
+                            text_parts.append(f"<assistant> {content}")
+                    
+                    if len(text_parts) >= 2:
+                        conversation_text = " ".join(text_parts) + " <|endoftext|>"
+                        self.conversations.append(conversation_text)               
+                except json.JSONDecodeError:
+                    continue    
+    def load_text(self):
+        with open(self.data_file, 'r', encoding='utf-8') as f:
+            content = f.read()
+            conversations = content.split('<|endoftext|>\n')
+            for conv in conversations:
+                conv = conv.strip()
+                if conv:
+                    self.conversations.append(conv + " <|endoftext|>")   
+    def get_tokenized_conversations(self, include_stats=False):
+        tokenized = []
+        stats = {'total_tokens': 0, 'truncated': 0, 'avg_length': 0}       
+        for conv in self.conversations:
+            tokens = self.tokenizer.encode_ids(conv)
+            if len(tokens) > self.max_length:
+                tokens = tokens[:self.max_length]
+                stats['truncated'] += 1
+            tokenized.append(tokens)
+            stats['total_tokens'] += len(tokens)        
+        if tokenized:
+            stats['avg_length'] = stats['total_tokens'] / len(tokenized)        
+        if include_stats:
+            return tokenized, stats
+        return tokenized    
+    def create_training_examples(self, stride: int = None):
+        if stride is None:
+            stride = self.max_length // 2        
+        examples = []
+        for conv in self.conversations:
+            tokens = self.tokenizer.encode_ids(conv)
+            if len(tokens) <= self.max_length:
+                examples.append(tokens)
+            else:
+                for i in range(0, len(tokens), stride):
+                    window = tokens[i:i + self.max_length]
+                    if len(window) >= 32:
+                        examples.append(window)
+        return examples
+def train_tokenizer_from_files(file_paths: List[str], 
+                              vocab_size: int = 32000,
+                              min_freq: int = 2,
+                              output_dir: str = "tokenizer",
+                              max_texts: int = None):
+    print(f"Training tokenizer with vocab_size={vocab_size}")
+    print(f"Input files: {file_paths}")    
+    all_texts = []
+    for file_path in file_paths:
+        print(f"Loading {file_path}...")        
+        if file_path.endswith('.jsonl'):
+            with open(file_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    try:
+                        conv = json.loads(line.strip())
+                        messages = conv.get("messages", [])
+                        text_parts = []
+                        for msg in messages:
+                            content = msg.get("content", "").strip()
+                            if content:
+                                text_parts.append(content)
+                        if text_parts:
+                            all_texts.append(" ".join(text_parts))
+                    except json.JSONDecodeError:
+                        continue
+        else:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+                chunks = content.split('\n\n')
+                for chunk in chunks:
+                    if chunk.strip():
+                        all_texts.append(chunk.strip())    
+    print(f"Loaded {len(all_texts)} texts")
+    if max_texts and len(all_texts) > max_texts:
+        import random
+        random.shuffle(all_texts)
+        all_texts = all_texts[:max_texts]
+        print(f"Limited to {len(all_texts)} texts")    
+    tokenizer = TechnicalTokenizer(vocab_size=vocab_size, min_freq=min_freq)
+    tokenizer.train_bpe(all_texts)
+    tokenizer.save(output_dir)    
+    print("\nTesting tokenization on sample texts:")
+    test_texts = [
+        "Hello, how can I help you with your Python programming question?",
+        "The neural network has 3 hidden layers with ReLU activation functions.",
+        "```python\ndef fibonacci(n):\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)\n```",
+        "The derivative of x^2 is 2x, and the integral is (x^3)/3 + C."
+    ]    
+    for text in test_texts:
+        tokenizer.analyze_tokenization(text)
+        print()    
+    return tokenizer
+def main():
+    parser = argparse.ArgumentParser(description="Train custom tokenizer for technical content")
+    parser.add_argument("--input_files", nargs='+', help="Input text/jsonl files")
+    parser.add_argument("--output_dir", default="tokenizer", help="Output directory for tokenizer")
+    parser.add_argument("--vocab_size", type=int, default=32000, help="Vocabulary size")
+    parser.add_argument("--min_freq", type=int, default=2, help="Minimum token frequency")
+    parser.add_argument("--max_texts", type=int, help="Maximum number of texts to use for training")
+    parser.add_argument("--test_file", help="Test file for analyzing tokenization")
+    parser.add_argument("--load_tokenizer", help="Load existing tokenizer from directory")    
+    args = parser.parse_args()    
+    default_input_file = "/kaggle/input/gpt-based-slm-dataset/slm_training_complete.jsonl"
+    default_text_file = "/kaggle/working/text_data/training_data_chat.txt"    
+    if not args.input_files and not args.load_tokenizer:
+        if os.path.exists(default_input_file):
+            args.input_files = [default_input_file]
+            print(f"No arguments provided, using default input file: {default_input_file}")
+        elif os.path.exists(default_text_file):
+            args.input_files = [default_text_file]
+            print(f"No arguments provided, using default text file: {default_text_file}")
+        else:
+            parser.error("No input files or tokenizer directory provided, and default files not found. "
+                         "Please specify --input_files or --load_tokenizer.")    
+    if args.load_tokenizer:
+        tokenizer = TechnicalTokenizer()
+        tokenizer.load(args.load_tokenizer)
+        if args.test_file:
+            print(f"\nTesting on {args.test_file}")
+            dataset = ConversationDataset(args.test_file, tokenizer)
+            tokenized, stats = dataset.get_tokenized_conversations(include_stats=True)
+            print(f"Dataset statistics:")
+            print(f"  Total conversations: {len(tokenized)}")
+            print(f"  Total tokens: {stats['total_tokens']:,}")
+            print(f"  Average tokens per conversation: {stats['avg_length']:.1f}")
+            print(f"  Conversations truncated: {stats['truncated']}")
+    else:
+        tokenizer = train_tokenizer_from_files(
+            file_paths=args.input_files,
+            vocab_size=args.vocab_size,
+            min_freq=args.min_freq,
+            output_dir=args.output_dir,
+            max_texts=args.max_texts
+        )
+        if args.test_file:
+            print(f"\nTesting on {args.test_file}")
+            dataset = ConversationDataset(args.test_file, tokenizer)
+            tokenized, stats = dataset.get_tokenized_conversations(include_stats=True)
+            print(f"Dataset statistics:")
+            print(f"  Total conversations: {len(tokenized)}")
+            print(f"  Total tokens: {stats['total_tokens']:,}")
+            print(f"  Average tokens per conversation: {stats['avg_length']:.1f}")
+            print(f"  Conversations truncated: {stats['truncated']}")
+
+if __name__ == "__main__":
+    main()