import torch import torch.nn as nn from typing import List, Dict, Optional import nltk from nltk.stem import WordNetLemmatizer, PorterStemmer import os import json import re from collections import Counter # Quietly ensure dependencies are present def ensure_nltk_data(): for pkg in ['punkt', 'wordnet', 'omw-1.4', 'punkt_tab']: try: nltk.data.find(f'tokenizers/{pkg}' if 'punkt' in pkg else f'corpora/{pkg}') except LookupError: nltk.download(pkg, quiet=True) ensure_nltk_data() # Optional SpaCy import try: import spacy SPACY_AVAILABLE = True except ImportError: SPACY_AVAILABLE = False print("SpaCy not available. Install with: pip install spacy") class AdvancedTokenizer: def __init__(self, vocab_size=5007, use_spacy_pos=False, use_stemming=False): self.vocab_size = vocab_size self.lemmatizer = WordNetLemmatizer() self.stemmer = PorterStemmer() self.use_spacy_pos = use_spacy_pos and SPACY_AVAILABLE self.use_stemming = use_stemming # Initialize SpaCy if requested if self.use_spacy_pos: try: self.nlp = spacy.load("en_core_web_sm") print("SpaCy POS Tagging: ENABLED") except OSError: print("SpaCy model not found. Disabling POS tagging.") self.use_spacy_pos = False self.word_to_id = {} self.id_to_word = {} # Special tokens self.unk_token = "" self.pad_token = "" self.sos_token = "" self.eos_token = "" self.instruction_token = "[INSTRUCTION]" self.response_token = "[RESPONSE]" self.system_token = "[SYSTEM]" self.user_token = "[USER]" self.thought_token = "[THOUGHT]" # Initialize with specials + Agentic CoT and Tool Calling Tokens self.specials = [ self.pad_token, self.sos_token, self.eos_token, self.unk_token, self.instruction_token, self.response_token, self.system_token, self.user_token, self.thought_token, "[TOOL_CALL]", "[TOOL_ARG]", "[TOOL_RESULT]" ] for i, token in enumerate(self.specials): self.word_to_id[token] = i self.id_to_word[i] = token self.is_trained = False def preprocess_text(self, text: str) -> str: """Applies basic grammar rules.""" # 1. Fix multiple spaces text = re.sub(r'\s+', ' ', text) # 2. Ensure space after punctuation text = re.sub(r'([.,;?!])(?=[a-zA-Z])', r'\1 ', text) # 3. Capitalize first letter of sentences sentences = re.split(r'([.?!]\s*)', text) text = ''.join([s.capitalize() if i % 2 == 0 else s for i, s in enumerate(sentences)]) return text.strip() def process_token(self, token: str, pos: Optional[str] = None) -> str: """Applies Lemmatization AND Stemming, optionally using POS tag.""" # 1. Lemmatize (with POS if available) if pos and pos in ['NOUN', 'VERB', 'ADJ', 'ADV']: # Map SpaCy POS to WordNet POS pos_map = {'NOUN': 'n', 'VERB': 'v', 'ADJ': 'a', 'ADV': 'r'} lemma = self.lemmatizer.lemmatize(token, pos=pos_map.get(pos, 'n')) else: lemma = self.lemmatizer.lemmatize(token) # 2. Stem (Root form) - OPTIONAL if self.use_stemming: lemma = self.stemmer.stem(lemma) return lemma def train(self, text: str): """Builds vocabulary from text.""" mode_str = "Grammar + Lemma + Stem" if self.use_spacy_pos: mode_str += " + POS" print(f"Training Tokenizer ({mode_str})...") pattern = '(' + '|'.join(map(re.escape, self.specials)) + ')' parts = re.split(pattern, text) processed_tokens = [] for part in parts: if part in self.specials: continue # Skip specials for vocab building elif part.strip(): part = self.preprocess_text(part) if self.use_spacy_pos: doc = self.nlp(part.lower()) processed_tokens.extend([self.process_token(token.text, token.pos_) for token in doc]) else: tokens = nltk.word_tokenize(part.lower()) processed_tokens.extend([self.process_token(t) for t in tokens]) # 3. Count frequencies counts = Counter(processed_tokens) # 4. Select top words num_specials = len(self.specials) most_common = counts.most_common(self.vocab_size - num_specials) for i, (word, _) in enumerate(most_common): idx = i + num_specials self.word_to_id[word] = idx self.id_to_word[idx] = word print(f"Tokenizer trained. Vocab size: {len(self.word_to_id)}") self.is_trained = True self.save("vocab.json") def encode(self, text: str) -> List[int]: if not self.is_trained: if os.path.exists("vocab.json"): self.load("vocab.json") else: self.train(text) ids = [] pattern = '(' + '|'.join(map(re.escape, self.specials)) + ')' parts = re.split(pattern, text) for part in parts: if part in self.specials: ids.append(self.word_to_id[part]) elif part.strip(): part = self.preprocess_text(part) if self.use_spacy_pos: doc = self.nlp(part.lower()) ids.extend([self.word_to_id.get(self.process_token(token.text, token.pos_), self.word_to_id[self.unk_token]) for token in doc]) else: tokens = nltk.word_tokenize(part.lower()) ids.extend([self.word_to_id.get(self.process_token(t), self.word_to_id[self.unk_token]) for t in tokens]) return ids def decode(self, ids: List[int]) -> str: tokens = [self.id_to_word.get(i, self.unk_token) for i in ids] # Simple detokenization (join with spaces) # NLTK doesn't have a perfect detokenizer built-in, simple join is okay for now return " ".join(tokens).replace(" .", ".").replace(" ,", ",") def save(self, path): with open(path, 'w') as f: json.dump(self.word_to_id, f) def load(self, path): if not os.path.exists(path): return with open(path, 'r') as f: self.word_to_id = json.load(f) self.id_to_word = {v: k for k, v in self.word_to_id.items()} self.is_trained = True # print(f"Tokenizer loaded. Vocab size: {len(self.word_to_id)}") # Singleton for easy access if needed, though class is better tokenizer = AdvancedTokenizer()