Create tokenizer.py

tokenizer.py

@@ -0,0 +1,142 @@
+import json
+import re
+from typing import List, Dict, Optional
+import numpy as np
+
+class VedaTokenizer:
+    """Custom tokenizer for Veda Programming LLM"""
+    
+    def __init__(self, vocab_size: int = 10000):
+        self.vocab_size = vocab_size
+        self.word_to_idx: Dict[str, int] = {}
+        self.idx_to_word: Dict[int, str] = {}
+        
+        # Special tokens
+        self.pad_token = "<PAD>"
+        self.unk_token = "<UNK>"
+        self.start_token = "<START>"
+        self.end_token = "<END>"
+        self.newline_token = "<NEWLINE>"
+        self.indent_token = "<INDENT>"
+        
+        self._init_special_tokens()
+    
+    def _init_special_tokens(self):
+        """Initialize special tokens"""
+        special_tokens = [
+            self.pad_token, 
+            self.unk_token, 
+            self.start_token, 
+            self.end_token,
+            self.newline_token,
+            self.indent_token
+        ]
+        for idx, token in enumerate(special_tokens):
+            self.word_to_idx[token] = idx
+            self.idx_to_word[idx] = token
+    
+    def _tokenize_code(self, text: str) -> List[str]:
+        """Tokenize code with special handling for programming constructs"""
+        # Replace newlines and indentation
+        text = text.replace('\n', f' {self.newline_token} ')
+        text = text.replace('\t', f' {self.indent_token} ')
+        text = text.replace('    ', f' {self.indent_token} ')
+        
+        # Tokenize with regex for code
+        pattern = r'''
+            \d+\.\d+|           # Floats
+            \d+|                # Integers
+            [a-zA-Z_]\w*|       # Identifiers
+            \"[^\"]*\"|         # Double quoted strings
+            \'[^\']*\'|         # Single quoted strings
+            \#[^\n]*|           # Comments
+            ==|!=|<=|>=|        # Comparison operators
+            \+=|-=|\*=|/=|      # Assignment operators
+            ->|=>|              # Arrow operators
+            \S                  # Other single characters
+        '''
+        tokens = re.findall(pattern, text, re.VERBOSE)
+        return tokens
+    
+    def fit(self, texts: List[str]):
+        """Build vocabulary from texts"""
+        word_freq = {}
+        
+        for text in texts:
+            tokens = self._tokenize_code(text)
+            for token in tokens:
+                word_freq[token] = word_freq.get(token, 0) + 1
+        
+        # Sort by frequency and take top vocab_size
+        sorted_words = sorted(word_freq.items(), key=lambda x: -x[1])
+        
+        start_idx = len(self.word_to_idx)
+        for idx, (word, _) in enumerate(sorted_words[:self.vocab_size - start_idx]):
+            actual_idx = idx + start_idx
+            self.word_to_idx[word] = actual_idx
+            self.idx_to_word[actual_idx] = word
+        
+        print(f"Vocabulary built with {len(self.word_to_idx)} tokens")
+    
+    def encode(self, text: str, max_length: Optional[int] = None) -> List[int]:
+        """Encode text to token indices"""
+        tokens = self._tokenize_code(text)
+        encoded = [self.word_to_idx.get(token, self.word_to_idx[self.unk_token]) 
+                   for token in tokens]
+        
+        if max_length:
+            if len(encoded) < max_length:
+                encoded += [self.word_to_idx[self.pad_token]] * (max_length - len(encoded))
+            else:
+                encoded = encoded[:max_length]
+        
+        return encoded
+    
+    def decode(self, indices: List[int]) -> str:
+        """Decode token indices back to text"""
+        tokens = []
+        for idx in indices:
+            if idx in self.idx_to_word:
+                token = self.idx_to_word[idx]
+                if token == self.pad_token:
+                    continue
+                elif token == self.newline_token:
+                    tokens.append('\n')
+                elif token == self.indent_token:
+                    tokens.append('    ')
+                else:
+                    tokens.append(token)
+        
+        # Join tokens intelligently
+        result = []
+        for i, token in enumerate(tokens):
+            if token in '.,;:)]}' or (i > 0 and tokens[i-1] in '([{'):
+                result.append(token)
+            elif token in '([{':
+                result.append(' ' + token if result else token)
+            else:
+                result.append(' ' + token if result else token)
+        
+        return ''.join(result).strip()
+    
+    def save(self, path: str):
+        """Save tokenizer to file"""
+        data = {
+            'vocab_size': self.vocab_size,
+            'word_to_idx': self.word_to_idx,
+            'idx_to_word': {str(k): v for k, v in self.idx_to_word.items()}
+        }
+        with open(path, 'w') as f:
+            json.dump(data, f)
+    
+    def load(self, path: str):
+        """Load tokenizer from file"""
+        with open(path, 'r') as f:
+            data = json.load(f)
+        self.vocab_size = data['vocab_size']
+        self.word_to_idx = data['word_to_idx']
+        self.idx_to_word = {int(k): v for k, v in data['idx_to_word'].items()}
+    
+    @property
+    def vocabulary_size(self) -> int:
+        return len(self.word_to_idx)