Create tokenizer_setup.py

tokenizer_setup.py

@@ -0,0 +1,120 @@
+import os
+import sentencepiece as spm
+from transformers import AutoTokenizer, PreTrainedTokenizerFast
+
+class TokenizerSetup:
+    def __init__(self, model_path="tokenizer", model_type="bpe", vocab_size=32000, hf_model=None):
+        """Initialize tokenizer setup for custom or pretrained use."""
+        self.model_path = model_path
+        self.model_type = model_type.lower()  # Normalize: bpe, unigram, char, word
+        self.vocab_size = vocab_size
+        self.hf_model = hf_model
+        self.tokenizer = None
+        
+        # Validate model_type
+        valid_types = ["bpe", "unigram", "char", "word"]
+        if self.model_type not in valid_types:
+            print(f"⚠️ Invalid model_type '{self.model_type}'. Choose from {valid_types}")
+            self.model_type = "bpe"
+
+    def train_sentencepiece(self, input_file):
+        """Train a SentencePiece tokenizer with specified settings."""
+        if not os.path.exists(input_file):
+            print(f"⚠️ Input file {input_file} not found! Provide a valid text corpus.")
+            return
+        
+        try:
+            spm.SentencePieceTrainer.Train(
+                f"--input={input_file} "
+                f"--model_prefix={self.model_path} "
+                f"--vocab_size={self.vocab_size} "
+                f"--model_type={self.model_type} "
+                f"--pad_id=0 --unk_id=1 --bos_id=2 --eos_id=3 "
+                f"--user_defined_symbols=<pad>,<unk>,<bos>,<eos>"  # Explicit special tokens
+            )
+            print(f"✅ Trained SentencePiece tokenizer. Saved as {self.model_path}.model")
+        except Exception as e:
+            print(f"⚠️ Error training SentencePiece: {e}")
+
+    def load_tokenizer(self):
+        """Load either a SentencePiece or Hugging Face tokenizer."""
+        try:
+            if self.hf_model:
+                self.tokenizer = AutoTokenizer.from_pretrained(self.hf_model)
+                print(f"✅ Loaded Hugging Face tokenizer from {self.hf_model}")
+            else:
+                sp_model = f"{self.model_path}.model"
+                if not os.path.exists(sp_model):
+                    print(f"⚠️ {sp_model} not found! Train it first.")
+                    return
+                
+                sp = spm.SentencePieceProcessor(model_file=sp_model)
+                self.tokenizer = PreTrainedTokenizerFast(
+                    tokenizer_object=sp,
+                    pad_token="<pad>",
+                    unk_token="<unk>",
+                    bos_token="<bos>",
+                    eos_token="<eos>"
+                )
+                print(f"✅ Loaded SentencePiece tokenizer from {sp_model}")
+        except Exception as e:
+            print(f"⚠️ Error loading tokenizer: {e}")
+
+    def save_tokenizer(self, save_dir="tokenizer/"):
+        """Save tokenizer files to a directory."""
+        if not self.tokenizer:
+            print("⚠️ No tokenizer loaded to save!")
+            return
+        
+        try:
+            os.makedirs(save_dir, exist_ok=True)
+            self.tokenizer.save_pretrained(save_dir)
+            if not self.hf_model:  # Copy SentencePiece files
+                for ext in [".model", ".vocab"]:
+                    src = f"{self.model_path}{ext}"
+                    if os.path.exists(src):
+                        os.system(f"cp {src} {save_dir}")
+            print(f"✅ Tokenizer saved to {save_dir}")
+        except Exception as e:
+            print(f"⚠️ Error saving tokenizer: {e}")
+
+    def tokenize_text(self, text, return_tensors=True):
+        """Tokenize text and show both IDs and decoded output."""
+        if not self.tokenizer:
+            print("⚠️ No tokenizer initialized! Load or train one first.")
+            return None
+        
+        try:
+            tokens = self.tokenizer(text, return_tensors="pt" if return_tensors else None)
+            ids = tokens["input_ids"] if return_tensors else tokens
+            decoded = self.tokenizer.decode(ids[0] if return_tensors else ids, skip_special_tokens=True)
+            print(f"🔹 Token IDs: {ids}")
+            print(f"🔹 Decoded: {decoded}")
+            return tokens
+        except Exception as e:
+            print(f"⚠️ Error tokenizing text: {e}")
+            return None
+
+if __name__ == "__main__":
+    # Setup with Charm 15 context
+    tokenizer_setup = TokenizerSetup(
+        model_path="tokenizer",
+        model_type="bpe",           # Matches your earlier BPE config
+        vocab_size=32000,           # Matches Mistral/Charm 15
+        hf_model=None               # Custom training; set to "mistralai/Mixtral-8x7B-Instruct-v0.1" for pretrained
+    )
+
+    # Train on Eclipse Corpuz (or other corpus)
+    input_file = "../datasets/eclipse_corpuz_1.1.txt"  # Adjust to your dataset
+    if not os.path.exists(f"{tokenizer_setup.model_path}.model"):
+        tokenizer_setup.train_sentencepiece(input_file)
+
+    # Load tokenizer
+    tokenizer_setup.load_tokenizer()
+
+    # Save for Charm 15 use
+    tokenizer_setup.save_tokenizer("../finetuned_charm15/")  # Match your training dir
+
+    # Test with sample
+    sample_text = "Charm 15 is an AI model optimized for deep learning and security."
+    tokenizer_setup.tokenize_text(sample_text)