Method evaluate

evals/ppl.py:114–166

Evaluate perplexity on the entire dataset

(self, dataset: Dataset)

    def evaluate(self, dataset: Dataset) -> Dict[str, Any]:
        """Evaluate perplexity on the entire dataset"""
        total_loss = 0
        total_tokens = 0
        total_sentences = 0
        
        # Initialize block statistics
        num_blocks = (self.block_size - 1) // self.bucket_size + 1
        block_loss = [torch.tensor(0., dtype=torch.float, device=self.device) for _ in range(num_blocks)]
        block_tokens = [1e-10 for _ in range(num_blocks)]
        bucket_sizes = [0 for _ in range(num_blocks)]
        
        # Create progress bar
        bar = tqdm(self.batchify(dataset, self.block_size))
        
        for batch in bar:
            batch_outputs = self.process_batch(batch)
            input_ids = batch_outputs['input_ids']
            
            nlls = batch_outputs['nlls']
            labels = batch_outputs['labels']
            blocks = batch_outputs['blocks']
            
            # Update statistics
            total_tokens += input_ids.ne(self.loss_fct.ignore_index).sum()
            total_sentences += input_ids.shape[0]
            print(input_ids.shape[1])

            for i in blocks:
                bucket_sizes[i] += 1
                
            # Calculate block-level loss
            for i, j in enumerate(range(0, min(input_ids.shape[-1], self.block_size), self.bucket_size)):
                block_loss[i] += nlls[:, j:j+self.bucket_size].sum()
                block_tokens[i] += labels[:, j:j+self.bucket_size].ne(self.loss_fct.ignore_index).sum()
            
            # Update total loss
            total_loss += batch_outputs['loss'].item() * labels.ne(self.loss_fct.ignore_index).sum()
            
            # Update progress bar
            ppls = [f"{math.exp(loss / toks):6.2f}" for loss, toks in zip(block_loss, block_tokens)]
            bar.set_description_str(f"[{total_tokens:10} tokens, {total_sentences:8} sentences] " + ' '.join(ppls))
        
        # Calculate final results
        final_ppl = math.exp(total_loss / total_tokens)
        block_ppls = [math.exp(loss / toks) for loss, toks in zip(block_loss, block_tokens)]
        
        return {
            'perplexity': final_ppl,
            'block_perplexities': block_ppls,
            'total_tokens': total_tokens,
            'total_sentences': total_sentences
        }