Last commit not found
| """ | |
| This script defines a custom tokenizer, `SupplyChainTokenizer`, specifically designed | |
| for a collaborative predictive supply chain model using Transformer-based | |
| architecture. It leverages a custom, industry-specific vocabulary (loaded from | |
| `vocab.json`) to prioritize domain-relevant tokens (SKUs, store IDs, plant IDs, | |
| promotion types, etc.) while employing Byte-Pair Encoding (BPE) to handle | |
| out-of-vocabulary words and variations. | |
| The script also includes a comprehensive example usage section demonstrating | |
| how to create, train, use, save, and load the tokenizer. This tokenizer is a | |
| critical component for bridging the gap between raw supply chain data and | |
| a Transformer-based forecasting model. | |
| """ | |
| import json | |
| import os | |
| from typing import List, Dict, Union, Tuple | |
| from tokenizers import ( | |
| Tokenizer, | |
| models, | |
| normalizers, | |
| pre_tokenizers, | |
| decoders, | |
| trainers, | |
| processors, | |
| ) | |
| from tokenizers.pre_tokenizers import WhitespaceSplit, Digits | |
| from tokenizers import Regex | |
| import pandas as pd | |
| class SupplyChainTokenizer: | |
| """ | |
| A custom tokenizer designed for the Enhanced Business Model for Collaborative | |
| Predictive Supply Chain. It prioritizes industry-specific tokens from a | |
| `vocab.json` file and uses Byte-Pair Encoding (BPE) for out-of-vocabulary | |
| (OOV) words. It handles various data types expected in supply chain data. | |
| Args: | |
| vocab_path (str): Path to the `vocab.json` file. | |
| max_length (int, optional): Maximum sequence length. Defaults to 512. | |
| """ | |
| def __init__(self, vocab_path: str, max_length: int = 512): | |
| if not os.path.exists(vocab_path): | |
| raise FileNotFoundError(f"Vocabulary file not found: {vocab_path}") | |
| self.vocab_path = vocab_path | |
| self.max_length = max_length | |
| # Load the custom vocabulary | |
| with open(self.vocab_path, "r", encoding="utf-8") as f: | |
| self.vocab = json.load(f) | |
| # 1. Create the BPE model | |
| self.bpe_model = models.BPE( | |
| vocab=self.vocab, # Initialize with the custom vocabulary | |
| merges=[], # We'll populate merges during training | |
| unk_token="[UNK]", # Unknown token | |
| ) | |
| # 2. Create a Tokenizer instance | |
| self.tokenizer = Tokenizer(self.bpe_model) | |
| # 3. Normalization (Lowercase and Unicode normalization) | |
| self.tokenizer.normalizer = normalizers.Sequence( | |
| [normalizers.NFD(), normalizers.Lowercase(), normalizers.StripAccents()] | |
| ) | |
| # 4. Pre-tokenization (Splitting into words) | |
| self.tokenizer.pre_tokenizer = pre_tokenizers.Sequence( | |
| [WhitespaceSplit(), Digits(individual_digits=True)] | |
| ) | |
| # 5. Decoder (Convert token IDs back to strings) | |
| self.tokenizer.decoder = decoders.BPEDecoder() | |
| # 6. Post-processing (Special tokens) | |
| self.tokenizer.post_processor = processors.TemplateProcessing( | |
| single="[CLS] $A [SEP]", | |
| pair="[CLS] $A [SEP] $B:1 [SEP]:1", | |
| special_tokens=[("[CLS]", self.vocab["[CLS]"]), ("[SEP]", self.vocab["[SEP]"])], | |
| ) | |
| # Adding this, although not used in encode or encode_as_ids | |
| self.pad_token_id = self.vocab["[PAD]"] | |
| def train_bpe(self, files: Union[str, List[str]], vocab_size: int = 30000): | |
| """ | |
| Trains the BPE model on text files. This updates the `merges` of the | |
| BPE model. This is *crucial* for handling words not in the initial | |
| `vocab.json`. | |
| Args: | |
| files (Union[str, List[str]]): Path(s) to text file(s) for training. | |
| vocab_size (int): The desired vocabulary size (including special tokens | |
| and initial vocabulary). | |
| """ | |
| if isinstance(files, str): | |
| files = [files] | |
| # Create a trainer | |
| trainer = trainers.BpeTrainer( | |
| vocab_size=vocab_size, | |
| special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"], | |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), # All single bytes | |
| show_progress=True, | |
| ) | |
| # Train the tokenizer | |
| self.tokenizer.train(files, trainer=trainer) | |
| def encode(self, text: str, text_pair: str = None) -> List[str]: | |
| """ | |
| Encodes text into a list of tokens. | |
| Args: | |
| text (str): The input text. | |
| text_pair (str, optional): An optional second input string. | |
| Returns: | |
| List[str]: A list of tokens. | |
| """ | |
| encoded = self.tokenizer.encode(text, text_pair) | |
| return encoded.tokens | |
| def encode_as_ids(self, text: str, text_pair: str = None) -> List[int]: | |
| """ | |
| Encodes text into a list of token IDs. | |
| Args: | |
| text (str): The input text. | |
| text_pair (str, optional): An optional second input string. | |
| Returns: | |
| List[int]: A list of token IDs. | |
| """ | |
| encoded = self.tokenizer.encode(text, text_pair) | |
| return encoded.ids | |
| def decode(self, ids: List[int], skip_special_tokens: bool = True) -> str: | |
| """ | |
| Decodes a list of token IDs back into a string. | |
| Args: | |
| ids (List[int]): The list of token IDs. | |
| skip_special_tokens (bool): Whether to skip special tokens in decoding. | |
| Returns: | |
| str: The decoded string. | |
| """ | |
| return self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens) | |
| def token_to_id(self, token: str) -> int: | |
| """ | |
| Converts a token to its corresponding ID. | |
| Args: | |
| token (str): The token. | |
| Returns: | |
| int: The token ID. Returns None if the token is not in the vocabulary. | |
| """ | |
| return self.vocab.get(token, self.vocab.get("[UNK]")) | |
| def id_to_token(self, id_: int) -> str: | |
| """ | |
| Converts a token ID to its corresponding token. | |
| Args: | |
| id_ (int): The token ID. | |
| Returns: | |
| str: The token. Returns "[UNK]" if the ID is not in the vocabulary. | |
| """ | |
| # Reverse lookup (efficient if needed frequently) | |
| reverse_vocab = {v: k for k, v in self.vocab.items()} | |
| return reverse_vocab.get(id_, "[UNK]") | |
| def get_vocab_size(self) -> int: | |
| """Gets the vocabulary size.""" | |
| return len(self.vocab) | |
| def save(self, directory: str, prefix: str = None): | |
| """ | |
| Saves the tokenizer configuration and vocabulary to a directory. | |
| Args: | |
| directory (str): The directory to save to. | |
| prefix (str, optional): An optional prefix for the filenames. | |
| """ | |
| if not os.path.exists(directory): | |
| os.makedirs(directory) | |
| # Save the tokenizer configuration | |
| self.tokenizer.save(os.path.join(directory, (prefix + "-" if prefix else "") + "tokenizer.json")) | |
| # Save a copy of the vocabulary (for easy access) | |
| with open(os.path.join(directory, (prefix + "-" if prefix else "") + "vocab.json"), "w", encoding="utf-8") as f: | |
| json.dump(self.vocab, f, ensure_ascii=False, indent=4) | |
| def from_pretrained(directory: str, prefix: str = None): | |
| """ | |
| Loads a pre-trained tokenizer from a directory. | |
| Args: | |
| directory (str): The directory to load from. | |
| prefix (str, optional): The optional prefix used when saving. | |
| Returns: | |
| SupplyChainTokenizer: The loaded tokenizer. | |
| """ | |
| vocab_path = os.path.join(directory, (prefix + "-" if prefix else "") + "vocab.json") | |
| # You could load the tokenizer.json, but since we have a custom class | |
| # with training logic, it's better to reconstruct the object this way. | |
| tokenizer = SupplyChainTokenizer(vocab_path) | |
| tokenizer.tokenizer = Tokenizer.from_file(os.path.join(directory, (prefix + "-" if prefix else "") + "tokenizer.json")) | |
| return tokenizer | |
| def prepare_for_model(self, data: pd.DataFrame) -> Tuple[List[List[int]], List[List[int]]]: | |
| """ | |
| Prepares a Pandas DataFrame for the Transformer model. This is the | |
| key method that integrates the tokenizer with the data. | |
| Args: | |
| data (pd.DataFrame): The input DataFrame, expected to have columns | |
| like 'timestamp', 'sku', 'store_id', 'quantity', 'price', | |
| 'discount', 'promotion_id', etc. The exact columns depend on | |
| the features you're using. | |
| Returns: | |
| Tuple[List[List[int]], List[List[int]]]: A tuple. | |
| 1. input_ids: List of token ID sequences for the model. | |
| 2. attention_mask: List of attention masks (1 for real tokens, 0 for padding). | |
| """ | |
| input_ids = [] | |
| attention_masks = [] | |
| for _, row in data.iterrows(): | |
| # Build the input string. This is where you define *how* your | |
| # features are combined into a single sequence. | |
| input_string = ( | |
| f"[CLS] timestamp: {row['timestamp']} " | |
| f"sku: {row['sku']} store_id: {row['store_id']} " | |
| f"quantity: {row['quantity']} price: {row['price']} " | |
| f"discount: {row['discount']} " | |
| ) | |
| # Add promotion information if available | |
| if 'promotion_id' in row and not pd.isna(row['promotion_id']): | |
| input_string += f"promotion_id: {row['promotion_id']} " | |
| # Add any other relevant features here | |
| if 'product_category' in row: | |
| input_string += f"product_category: {row['product_category']} " | |
| # Add other external features | |
| input_string += "[SEP]" | |
| # Tokenize | |
| encoded = self.tokenizer.encode(input_string) | |
| token_ids = encoded.ids | |
| attention_mask = encoded.attention_mask | |
| # Padding (up to max_length) | |
| padding_length = self.max_length - len(token_ids) | |
| if padding_length > 0: | |
| token_ids += [self.pad_token_id] * padding_length | |
| attention_mask += [0] * padding_length | |
| elif padding_length < 0: # Truncation | |
| token_ids = token_ids[:self.max_length] | |
| attention_mask = attention_mask[:self.max_length] | |
| input_ids.append(token_ids) | |
| attention_masks.append(attention_mask) | |
| return input_ids, attention_masks | |
| # Example Usage (Illustrative) | |
| if __name__ == "__main__": | |
| # --- Create a dummy vocab.json --- | |
| vocab = { | |
| "[UNK]": 0, | |
| "[CLS]": 1, | |
| "[SEP]": 2, | |
| "[PAD]": 3, | |
| "[MASK]": 4, | |
| "timestamp:": 5, | |
| "sku:": 6, | |
| "store_id:": 7, | |
| "quantity:": 8, | |
| "price:": 9, | |
| "discount:": 10, | |
| "promotion_id:": 11, | |
| "product_category:": 12, | |
| "SKU123": 13, # Example SKU | |
| "SKU123-RED": 14, # Example SKU variant | |
| "SKU123-BLUE": 15, | |
| "STORE456": 16, # Example store ID | |
| "PLANT789": 17, # Example plant ID | |
| "WHOLESALER001": 18, # Example Wholesaler | |
| "RETAILER002": 19, # Example Retailer | |
| "BOGO": 20, | |
| "DISCOUNT":21, | |
| } | |
| with open("vocab.json", "w") as f: | |
| json.dump(vocab, f, indent=4) | |
| # --- Create the tokenizer --- | |
| tokenizer = SupplyChainTokenizer(vocab_path="vocab.json") | |
| # --- Example training (on a dummy text file) --- | |
| with open("training_data.txt", "w", encoding="utf-8") as f: | |
| f.write("This is some example text for training the BPE model.\n") | |
| f.write("SKU123 is a product. STORE456 is another. plant789 is, too.\n") | |
| f.write("This file contains words not in the initial vocabulary.\n") | |
| tokenizer.train_bpe("training_data.txt", vocab_size=50) # Small vocab for the example | |
| # --- Example encoding --- | |
| text = "timestamp: 2024-07-03 sku: SKU123 store_id: STORE456 quantity: 2 price: 10.99 discount: 0.0" | |
| encoded_tokens = tokenizer.encode(text) | |
| encoded_ids = tokenizer.encode_as_ids(text) | |
| print(f"Encoded tokens: {encoded_tokens}") | |
| print(f"Encoded IDs: {encoded_ids}") | |
| decoded_text = tokenizer.decode(encoded_ids) | |
| print(f"Decoded text: {decoded_text}") | |
| # -- Example with DataFrame --- | |
| data = { | |
| 'timestamp': ['2024-07-03 10:00:00', '2024-07-03 11:00:00'], | |
| 'sku': ['SKU123', 'SKU123-RED'], | |
| 'store_id': ['STORE456', 'STORE456'], | |
| 'quantity': [2, 1], | |
| 'price': [10.99, 12.99], | |
| 'discount': [0.0, 1.0], | |
| 'promotion_id': ['BOGO', None], | |
| 'product_category': ['Electronics', 'Electronics'] | |
| } | |
| df = pd.DataFrame(data) | |
| input_ids, attention_masks = tokenizer.prepare_for_model(df) | |
| print(f"Input IDs (for model): {input_ids}") | |
| print(f"Attention Masks: {attention_masks}") | |
| # --- Save and load --- | |
| tokenizer.save("my_tokenizer") | |
| loaded_tokenizer = SupplyChainTokenizer.from_pretrained("my_tokenizer") | |
| print(f"Loaded tokenizer vocab size: {loaded_tokenizer.get_vocab_size()}") | |
| # Clean up example files | |
| os.remove("vocab.json") | |
| os.remove("training_data.txt") |