Update tokenizer.py

tokenizer.py

@@ -1,126 +1,98 @@
 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):
+    def __init__(self, vocab_size: int = 5000):
         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
-        ]
+        special_tokens = ["<PAD>", "<UNK>", "<START>", "<END>", "<NL>", "<INDENT>"]
         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 code"""
+        text = text.replace('\n', ' <NL> ')
+        text = text.replace('\t', ' <INDENT> ')
+        text = text.replace('    ', ' <INDENT> ')
         
-        # 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
+        pattern = r'[a-zA-Z_]\w*|[0-9]+\.?[0-9]*|\"[^\"]*\"|\'[^\']*\'|==|!=|<=|>=|[^\s]'
+        tokens = re.findall(pattern, text)
+        return [t for t in tokens if t.strip()]
     
     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:
+            for token in self._tokenize_code(text):
                 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")
+        print(f"Vocabulary: {len(self.word_to_idx)} tokens")
     
     def encode(self, text: str, max_length: Optional[int] = None) -> List[int]:
-        """Encode text to token indices"""
+        """Encode text to indices"""
         tokens = self._tokenize_code(text)
-        encoded = [self.word_to_idx.get(token, self.word_to_idx[self.unk_token]) 
-                   for token in tokens]
+        encoded = [self.word_to_idx.get(t, 1) for t in tokens]  # 1 = UNK
         
         if max_length:
             if len(encoded) < max_length:
-                encoded += [self.word_to_idx[self.pad_token]] * (max_length - len(encoded))
+                encoded += [0] * (max_length - len(encoded))  # 0 = PAD
             else:
                 encoded = encoded[:max_length]
-        
         return encoded
     
     def decode(self, indices: List[int]) -> str:
-        """Decode token indices back to text"""
+        """Decode indices 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:
+                if token == "<PAD>":
                     continue
-                elif token == self.newline_token:
+                elif token == "<NL>":
                     tokens.append('\n')
-                elif token == self.indent_token:
+                elif token == "<INDENT>":
                     tokens.append('    ')
+                elif token in ["<UNK>", "<START>", "<END>"]:
+                    continue
                 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 '([{'):
+            if token in '\n':
+                result.append(token)
+            elif token == '    ':
+                result.append(token)
+            elif token in '.,;:)]}':
+                result.append(token)
+            elif i > 0 and tokens[i-1] in '([{':
                 result.append(token)
             elif token in '([{':
-                result.append(' ' + token if result else token)
+                result.append(' ' + token if result and result[-1] not in '\n ' else token)
             else:
-                result.append(' ' + token if result else token)
+                result.append(' ' + token if result and result[-1] not in '\n    ' else token)
         
         return ''.join(result).strip()
     
     def save(self, path: str):
-        """Save tokenizer to file"""
+        """Save tokenizer"""
         data = {
             'vocab_size': self.vocab_size,
             'word_to_idx': self.word_to_idx,
@@ -130,7 +102,7 @@ class VedaTokenizer:
             json.dump(data, f)
     
     def load(self, path: str):
-        """Load tokenizer from file"""
+        """Load tokenizer"""
         with open(path, 'r') as f:
             data = json.load(f)
         self.vocab_size = data['vocab_size']