""" PyTorch Dataset for anime filename token classification. Loads JSONL data (tokens + BIO labels) and converts to model inputs. Handles token-ID conversion, label encoding, padding, and truncation. """ import json from collections import Counter import numpy as np import torch from torch.utils.data import Dataset from typing import Dict, List, Optional, Sequence, Tuple from .config import Config from .label_repairs import repair_sequel_season_labels from .tokenizer import AnimeTokenizer def encode_token_classification_values( item: Dict, tokenizer: AnimeTokenizer, label2id: Dict[str, int], max_length: int, apply_label_repairs: bool = True, vocab: Optional[Dict[str, int]] = None, ) -> Tuple[List[int], List[bool], List[int]]: tokens, labels = training_labels_for_tokenizer(item, tokenizer, apply_label_repairs) token_vocab = vocab if vocab is not None else tokenizer.get_vocab() unk_id = tokenizer.unk_token_id if tokenizer.unk_token_id is not None else 1 input_ids = [token_vocab.get(token, unk_id) for token in tokens] input_ids = [tokenizer.cls_token_id] + input_ids + [tokenizer.sep_token_id] label_ids: List[int] = [-100] label_ids.extend(label2id.get(label, 0) for label in labels) label_ids.append(-100) attention_mask = [1] * len(input_ids) if len(input_ids) > max_length: input_ids = [input_ids[0]] + input_ids[1:max_length - 1] + [input_ids[-1]] label_ids = [label_ids[0]] + label_ids[1:max_length - 1] + [label_ids[-1]] attention_mask = [attention_mask[0]] + attention_mask[1:max_length - 1] + [attention_mask[-1]] pad_len = max_length - len(input_ids) if pad_len > 0: input_ids += [tokenizer.pad_token_id] * pad_len label_ids += [-100] * pad_len attention_mask += [0] * pad_len return input_ids, [bool(value) for value in attention_mask], label_ids def encode_token_classification_item( item: Dict, tokenizer: AnimeTokenizer, label2id: Dict[str, int], max_length: int, apply_label_repairs: bool = True, vocab: Optional[Dict[str, int]] = None, ) -> Dict[str, torch.Tensor]: input_ids, attention_mask, label_ids = encode_token_classification_values( item, tokenizer, label2id, max_length, apply_label_repairs, vocab, ) return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "attention_mask": torch.tensor(attention_mask, dtype=torch.bool), "labels": torch.tensor(label_ids, dtype=torch.long), } class AnimeItemsDataset(Dataset): """Map-style dataset backed by already-loaded JSONL items.""" def __init__( self, data: Sequence[Dict], tokenizer: AnimeTokenizer, label2id: Dict[str, int], max_length: int = 64, apply_label_repairs: bool = True, ): self.data = data self.tokenizer = tokenizer self.label2id = label2id self.max_length = max_length self.apply_label_repairs = apply_label_repairs self.vocab = tokenizer.get_vocab() def __len__(self) -> int: return len(self.data) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: return encode_token_classification_item( self.data[idx], self.tokenizer, self.label2id, self.max_length, self.apply_label_repairs, self.vocab, ) class EncodedAnimeDataset(Dataset): """Dataset that stores padded tensors so training workers do no token work.""" def __init__( self, data: Sequence[Dict], tokenizer: AnimeTokenizer, label2id: Dict[str, int], max_length: int = 64, device: Optional[torch.device] = None, apply_label_repairs: bool = True, ): target_device = device or torch.device("cpu") vocab = tokenizer.get_vocab() input_ids = np.full( (len(data), max_length), tokenizer.pad_token_id, dtype=np.int64, ) attention_mask = np.zeros((len(data), max_length), dtype=np.bool_) labels = np.full((len(data), max_length), -100, dtype=np.int64) for idx, item in enumerate(data): item_input_ids, item_attention_mask, item_labels = encode_token_classification_values( item, tokenizer, label2id, max_length, apply_label_repairs, vocab, ) input_ids[idx] = item_input_ids attention_mask[idx] = item_attention_mask labels[idx] = item_labels self.input_ids = torch.from_numpy(input_ids).to(target_device) self.attention_mask = torch.from_numpy(attention_mask).to(target_device) self.labels = torch.from_numpy(labels).to(target_device) def __len__(self) -> int: return self.input_ids.shape[0] def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: return { "input_ids": self.input_ids[idx], "attention_mask": self.attention_mask[idx], "labels": self.labels[idx], } class AnimeDataset(AnimeItemsDataset): """ Dataset for anime filename token classification. Loads pre-tokenized data from JSONL files and prepares model inputs. Each sample has: - input_ids: token IDs with [CLS] prefix and [SEP] suffix - attention_mask: 1 for real tokens, 0 for padding - labels: integer label IDs, -100 for special/padding tokens """ def __init__( self, data_path: str, tokenizer: AnimeTokenizer, label2id: Dict[str, int], max_length: int = 64, ): """ Args: data_path: Path to JSONL file with tokens and labels. tokenizer: AnimeTokenizer instance. label2id: Mapping from label string to integer ID. max_length: Maximum sequence length (including special tokens). """ data: List[Dict] = [] with open(data_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if line: data.append(json.loads(line)) super().__init__(data, tokenizer, label2id, max_length) def align_tokens_for_tokenizer( tokens: List[str], labels: List[str], tokenizer: AnimeTokenizer, ) -> tuple[List[str], List[str]]: """ Align pre-labeled JSONL samples to the selected tokenizer. The existing datasets store regex-tokenized samples. For the char A/B run, each original token is split into characters while preserving BIO spans: B-X stays on the first character, and the rest become I-X. """ if getattr(tokenizer, "tokenizer_variant", "regex") != "char": return tokens, labels aligned_tokens: List[str] = [] aligned_labels: List[str] = [] for token, label in zip(tokens, labels): pieces = tokenizer.tokenize(token) if not pieces: continue aligned_tokens.extend(pieces) aligned_labels.append(label) if label.startswith(("B-", "I-")): continuation = "I-" + label.split("-", 1)[1] else: continuation = label aligned_labels.extend([continuation] * (len(pieces) - 1)) return aligned_tokens, aligned_labels def labels_for_tokenizer( item: Dict, tokenizer: AnimeTokenizer, ) -> Tuple[List[str], List[str]]: """ Return tokens and labels in the exact tokenizer space used by the model. Older DMHY weak-label files store a post-processed token sequence where group/title brackets may be expanded even though AnimeTokenizer keeps the same bracketed text as one inference token. If the raw filename is present, project those weak labels back to character spans and then onto the current tokenizer output. This keeps train/eval/inference preprocessing identical. """ filename = item.get("filename") source_tokens, source_labels, _repairs = repair_sequel_season_labels(item) tokenizer_variant = getattr(tokenizer, "tokenizer_variant", "regex") if not filename: return align_tokens_for_tokenizer(source_tokens, source_labels, tokenizer) # Current char datasets are already in the exact inference token space. # Avoid re-scanning every filename during training. if item.get("tokenizer_variant") == tokenizer_variant: target_tokens = tokenizer.tokenize(filename) if source_tokens == target_tokens: return source_tokens, source_labels projected = project_labels_from_filename( filename=filename, source_tokens=source_tokens, source_labels=source_labels, tokenizer=tokenizer, ) if projected is not None: return projected # Fall back to the legacy behavior for synthetic fixtures or malformed rows. return align_tokens_for_tokenizer(source_tokens, source_labels, tokenizer) def training_labels_for_tokenizer( item: Dict, tokenizer: AnimeTokenizer, apply_label_repairs: bool, ) -> Tuple[List[str], List[str]]: """Fast path for authoritative char JSONL rows used in full training.""" tokenizer_variant = getattr(tokenizer, "tokenizer_variant", "regex") if not apply_label_repairs and item.get("tokenizer_variant") == tokenizer_variant: tokens = item.get("tokens", []) labels = item.get("labels", []) filename = item.get("filename") if len(tokens) == len(labels): if tokenizer_variant != "char" or filename is None or tokens == list(str(filename)): return tokens, labels return labels_for_tokenizer(item, tokenizer) def token_offsets_in_text(text: str, tokens: List[str]) -> Optional[List[Tuple[int, int]]]: """Find token character offsets by scanning left to right.""" offsets: List[Tuple[int, int]] = [] cursor = 0 for token in tokens: if token == "": offsets.append((cursor, cursor)) continue start = text.find(token, cursor) if start < 0: return None end = start + len(token) offsets.append((start, end)) cursor = end return offsets def project_source_labels_to_chars( text: str, source_tokens: List[str], source_labels: List[str], ) -> Optional[List[str]]: """Project source token BIO labels to per-character entity names.""" offsets = token_offsets_in_text(text, source_tokens) if offsets is None or len(source_tokens) != len(source_labels): return None char_entities = ["O"] * len(text) for token, label, (start, end) in zip(source_tokens, source_labels, offsets): if not label.startswith(("B-", "I-")): continue entity = label.split("-", 1)[1] # Bracketed single-token metadata in older data often includes the # brackets in the token. Keep container punctuation as O so a tokenizer # that splits brackets can learn cleaner boundaries. inner_start = start inner_end = end if len(token) >= 2 and token[0] in "[【(《" and token[-1] in "]】)》": inner_start += 1 inner_end -= 1 for pos in range(inner_start, inner_end): if 0 <= pos < len(char_entities): char_entities[pos] = entity return char_entities def labels_from_char_projection( text: str, target_tokens: List[str], char_entities: List[str], ) -> Optional[List[str]]: """Assign legal IOB2 labels to target tokens from per-character entities.""" offsets = token_offsets_in_text(text, target_tokens) if offsets is None: return None labels: List[str] = [] active_entity: Optional[str] = None for start, end in offsets: span_entities = [ char_entities[pos] for pos in range(start, end) if 0 <= pos < len(char_entities) and char_entities[pos] != "O" ] if not span_entities: labels.append("O") active_entity = None continue entity = Counter(span_entities).most_common(1)[0][0] prefix = "I" if active_entity == entity else "B" labels.append(f"{prefix}-{entity}") active_entity = entity return labels def project_labels_from_filename( filename: str, source_tokens: List[str], source_labels: List[str], tokenizer: AnimeTokenizer, ) -> Optional[Tuple[List[str], List[str]]]: """ Re-tokenize filename and project weak BIO labels onto that tokenizer. Returns None when source tokens cannot be aligned to the filename. """ char_entities = project_source_labels_to_chars(filename, source_tokens, source_labels) if char_entities is None: return None target_tokens = tokenizer.tokenize(filename) target_labels = labels_from_char_projection(filename, target_tokens, char_entities) if target_labels is None or len(target_tokens) != len(target_labels): return None return target_tokens, target_labels def create_datasets( data_path: str, tokenizer: AnimeTokenizer, config: Config, ) -> tuple: """ Create train and validation datasets from a JSONL file. The file is split by the first N samples for training, the rest for validation based on config.train_split. Returns: (train_dataset, eval_dataset) """ # Load all data to determine split with open(data_path, 'r', encoding='utf-8') as f: all_data = [json.loads(line) for line in f if line.strip()] split_idx = int(len(all_data) * config.train_split) train_data = all_data[:split_idx] eval_data = all_data[split_idx:] # Write temp files for each split import tempfile import os train_file = os.path.join(tempfile.gettempdir(), "anime_train.jsonl") eval_file = os.path.join(tempfile.gettempdir(), "anime_eval.jsonl") with open(train_file, 'w', encoding='utf-8') as f: for item in train_data: f.write(json.dumps(item, ensure_ascii=False) + '\n') with open(eval_file, 'w', encoding='utf-8') as f: for item in eval_data: f.write(json.dumps(item, ensure_ascii=False) + '\n') train_dataset = AnimeDataset( data_path=train_file, tokenizer=tokenizer, label2id=config.label2id, max_length=config.max_seq_length, ) eval_dataset = AnimeDataset( data_path=eval_file, tokenizer=tokenizer, label2id=config.label2id, max_length=config.max_seq_length, ) return train_dataset, eval_dataset if __name__ == "__main__": # Quick test cfg = Config() tok = AnimeTokenizer() # Build a minimal vocab tok.build_vocab([["[ANi]", "test", "S2", "-", "03"], ["[Baha]", "anime", "01"]]) ds = AnimeDataset( data_path="data/synthetic.jsonl", tokenizer=tok, label2id=cfg.label2id, max_length=cfg.max_seq_length, ) print(f"Dataset size: {len(ds)}") if len(ds) > 0: sample = ds[0] print(f"input_ids shape: {sample['input_ids'].shape}") print(f"attention_mask shape: {sample['attention_mask'].shape}") print(f"labels shape: {sample['labels'].shape}") print(f"input_ids: {sample['input_ids'].tolist()}") print(f"labels: {sample['labels'].tolist()}") print(f"attention_mask: {sample['attention_mask'].tolist()}")