sail_scripts/model/tokenizer.py

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 = "<UNK>"
        self.pad_token = "<PAD>"
        self.sos_token = "<SOS>"
        self.eos_token = "<EOS>"
        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()