Function create_compute_metrics_function

utils/compute_utils.py:3–90

Create a compute_metrics function with the required parameters.

(tokenizer, args, metric)

import numpy as np

def create_compute_metrics_function(tokenizer, args, metric):
    """Create a compute_metrics function with the required parameters."""
    
    task_to_labels = {
        'cola': ['unacceptable', 'acceptable'],
        'sst2': ['negative', 'positive'],
        'mrpc': ['not_equivalent', 'equivalent'],
        'qqp': ['not_duplicate', 'duplicate'],
        'stsb': None,  # Regression task
        'mnli': ['entailment', 'neutral', 'contradiction'],
        'qnli': ['entailment', 'not_entailment'],
        'rte': ['entailment', 'not_entailment'],
        'wnli': ['not_entailment', 'entailment']
    }
    
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        
        # Convert predictions to numpy array and get the first element if it's a tuple
        if isinstance(preds, tuple):
            preds = preds[0]
        
        # Ensure preds is a numpy array
        preds = np.array(preds)
        
        # Handle multi-dimensional arrays (e.g., from beam search)
        if len(preds.shape) > 2:
            preds = preds.reshape(-1, preds.shape[-1])
        
        # Convert to list of lists if necessary
        if isinstance(preds[0], np.ndarray):
            preds = [pred.tolist() for pred in preds]
        
        try:
            # Decode predictions
            decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
            # Clean up decoded predictions (remove extra whitespace)
            decoded_preds = [pred.strip() for pred in decoded_preds]
            
            if args.task == 'stsb':
                # Convert string predictions to float for regression task
                # Handle potential formatting issues
                decoded_preds = [float(pred.replace(',', '.')) if pred.strip() else 0.0 
                               for pred in decoded_preds]
            else:
                # Convert string predictions to label indices for classification tasks
                label_list = task_to_labels[args.task]
                decoded_preds = [label_list.index(pred) if pred in label_list 
                               else 0 for pred in decoded_preds]

            # Handle labels
            if labels is not None:
                # Replace -100 with pad token
                labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
                
                # Convert labels to list of lists if necessary
                if isinstance(labels[0], np.ndarray):
                    labels = [label.tolist() for label in labels]