code/data/indirect_idx/tokenizer.py

import string
from typing import List

class CharacterTokenizer:
    """Character-level tokenizer for indirect indexing task."""
    
    def __init__(self):
        self.chars = list(string.ascii_uppercase + string.ascii_lowercase + string.digits)
        
        # Special tokens
        self.PAD_TOKEN = '<PAD>'
        self.SEP_TOKEN = '<SEP>'  # For separating sequence, source, shift
        self.PLUS_TOKEN = '<PLUS>'  # For positive shifts
        self.MINUS_TOKEN = '<MINUS>'  # For negative shifts
        
        # Build vocabulary
        self.vocab = [self.PAD_TOKEN, self.SEP_TOKEN, self.PLUS_TOKEN, self.MINUS_TOKEN] + self.chars
        # Create mappings
        self.char_to_idx = {char: idx for idx, char in enumerate(self.vocab)}
        self.idx_to_char = {idx: char for idx, char in enumerate(self.vocab)}
        
        # Special token indices
        self.pad_idx = self.char_to_idx[self.PAD_TOKEN]
        self.sep_idx = self.char_to_idx[self.SEP_TOKEN]
        self.plus_idx = self.char_to_idx[self.PLUS_TOKEN]
        self.minus_idx = self.char_to_idx[self.MINUS_TOKEN]
        
    @property
    def vocab_size(self) -> int:
        return len(self.vocab)
    
    def encode_input(self, input_str: str) -> List[int]:
        """
        Format: "sequence, source_char, shift" -> [seq_tokens, SEP, source_token, SEP, shift_tokens]
        """
        parts = input_str.split(', ')
        if len(parts) != 3:
            raise ValueError(f"Invalid input format: {input_str}")
        
        sequence, source_char, shift_str = parts
        tokens = []
        
        # Encode sequence
        for char in sequence:
            tokens.append(self.char_to_idx.get(char))
        
        # Add separator
        tokens.append(self.sep_idx)
        
        # Encode source character
        tokens.append(self.char_to_idx.get(source_char))
        
        # Add separator
        tokens.append(self.sep_idx)
        
        # Encode shift (handle + and - signs)
        shift = int(shift_str)
        if shift >= 0:
            tokens.append(self.plus_idx)
            shift_digits = str(shift)
        else:
            tokens.append(self.minus_idx)
            shift_digits = str(abs(shift))
        
        # Add shift digits
        for digit in shift_digits:
            tokens.append(self.char_to_idx.get(digit))
        
        # Add separator
        tokens.append(self.sep_idx)
        return tokens
    
    def encode_target(self, target_char: str) -> int:
        return self.char_to_idx.get(target_char)
    
    def decode_tokens(self, tokens: List[int]) -> str:
        chars = []
        for token in tokens:
            if token == self.pad_idx:
                break  # Stop at padding
            chars.append(self.idx_to_char.get(token))
        return ''.join(chars)
    
    def decode_target(self, target_idx: int) -> str:
        return self.idx_to_char.get(target_idx)