| 14 | from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss |
| 15 | |
| 16 | class PerplexityEvaluator: |
| 17 | def __init__( |
| 18 | self, |
| 19 | model: PreTrainedModel, |
| 20 | tokenizer: PreTrainedTokenizer, |
| 21 | device: str = "cuda", |
| 22 | block_size: int = 32768, |
| 23 | bucket_size: int = 2048, |
| 24 | batch_size: int = 1 |
| 25 | ): |
| 26 | self.model = model |
| 27 | self.tokenizer = tokenizer |
| 28 | self.device = device |
| 29 | self.block_size = block_size |
| 30 | self.bucket_size = bucket_size |
| 31 | self.batch_size = batch_size |
| 32 | self.loss_fct = FusedCrossEntropyLoss(reduction='sum') |
| 33 | |
| 34 | |
| 35 | @staticmethod |
| 36 | def preprocess( |
| 37 | examples: Dict[str, List[Any]], |
| 38 | tokenizer: PreTrainedTokenizer, |
| 39 | column_name: str = 'text' |
| 40 | ) -> Dict[str, List[List[int]]]: |
| 41 | """Preprocess text data""" |
| 42 | tokenized = tokenizer(examples[column_name]) |
| 43 | return { |
| 44 | 'input_ids': tokenized['input_ids'], |
| 45 | 'length': [len(ids) for ids in tokenized['input_ids']] |
| 46 | } |
| 47 | |
| 48 | def batchify(self, dataset: Dataset, tokens_per_batch: int) -> Iterator[List[torch.Tensor]]: |
| 49 | """Split dataset into batches of exactly block_size length""" |
| 50 | current_tokens = [] # Buffer to store all tokens |
| 51 | |
| 52 | for sentence in dataset: |
| 53 | # Convert input_ids to list and add to buffer |
| 54 | tokens = sentence['input_ids'].tolist() if torch.is_tensor(sentence['input_ids']) else list(sentence['input_ids']) |
| 55 | if not tokens: |
| 56 | continue |
| 57 | current_tokens.extend(tokens) |
| 58 | |
| 59 | # When we have enough tokens, yield batches |
| 60 | while len(current_tokens) >= self.block_size * self.batch_size: |
| 61 | batch = [] |
| 62 | for _ in range(self.batch_size): |
| 63 | # Extract exactly block_size tokens |
| 64 | batch.append(torch.tensor(current_tokens[:self.block_size], dtype=torch.long)) |
| 65 | current_tokens = current_tokens[self.block_size:] |
| 66 | yield batch |
| 67 | |
| 68 | # Handle remaining tokens if they form complete blocks |
| 69 | if len(current_tokens) >= self.block_size: |
| 70 | remaining_batches = len(current_tokens) // self.block_size |
| 71 | remaining_batches = min(remaining_batches, self.batch_size) |
| 72 | if remaining_batches > 0: |
| 73 | batch = [] |