utils.py

from tqdm import tqdm

import io
import json
import multiprocessing
import os
import sys
import torch


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.cuda.manual_seed_all(seed)
    torch.manual_seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    
SEED = 42
set_seed(SEED)

def read_jsonl(file_path):
    data = []
    with open(file_path, 'r') as f:
        for line in f:
            data.append(json.loads(line))
    return data

def execute_code(code, output_queue):
    """Executes the given code and puts the result in the queue."""
    stdout_capture = io.StringIO()
    original_stdout = sys.stdout  # Save original stdout
    try:
        sys.stdout = stdout_capture  # Redirect stdout
        exec(code)  # Execute the code
        output_queue.put((True, stdout_capture.getvalue()))
    except Exception as e:
        output_queue.put((False, str(e)))
    finally:
        sys.stdout = original_stdout  # Always restore stdout

def get_python_code_output(code, timeout=3):
    """
    Runs Python code and captures output, skipping code that waits for stdin.
    Times out after the specified duration.
    """
    output_queue = multiprocessing.Queue()
    process = multiprocessing.Process(target=execute_code, args=(code, output_queue))
    
    process.start()
    process.join(timeout)

    # If process is still alive after timeout, terminate it
    if process.is_alive():
        process.terminate()  # Kill the process
        process.join()  # Ensure cleanup
        return None  # Skipping code due to timeout

    # Retrieve results from the queue
    if not output_queue.empty():
        success, result = output_queue.get()
        return result if success else None

    return None
    
def extract_code_block(text, start_delimiter="```python", end_delimiter="```"):
    """Extracts text between start_delimiter and end_delimiter."""
    try:
        start_index = text.index(start_delimiter) + len(start_delimiter)
        end_index = text.index(end_delimiter, start_index)
        return text[start_index:end_index].strip()
    except ValueError:
        return None

def safe_parse_float(value):
    try:
        return float(value)
    except (ValueError, TypeError):
        return None

def calculate_conditional_perplexity(prefix, generation, model, tokenizer):
    """
    Calculate the conditional perplexity of a generation given its prefix.
    
    Args:
        prefix (str): The prompt/prefix text used to generate the completion
        generation (str): The generated completion text
        model: The language model
        tokenizer: The tokenizer associated with the model
        
    Returns:
        float: The conditional perplexity score
    """
    # Tokenize the full sequence (prefix + generation)
    full_text = prefix + generation
    encodings = tokenizer(full_text, return_tensors="pt")
    input_ids = encodings["input_ids"].to("cuda")
    
    # Tokenize just the prefix to find its length
    prefix_encodings = tokenizer(prefix, return_tensors="pt")
    prefix_length = prefix_encodings["input_ids"].shape[1]
    
    # Create labels: -100 for prefix tokens (ignored in loss calculation) 
    # and actual token ids for generation
    labels = input_ids.clone()
    labels[:, :prefix_length] = -100
    
    # Calculate log-likelihood
    with torch.no_grad():
        outputs = model(input_ids=input_ids, labels=labels)
        
    neg_log_likelihood = outputs.loss.item()
    
    # Calculate perplexity only on the generated portion
    ppl = torch.exp(torch.tensor(neg_log_likelihood)).item()
    
    return ppl

def evaluate_generations(prefixes, generations, model, tokenizer):
    """
    Calculate conditional perplexity for a list of prefix-generation pairs.
    
    Args:
        prefixes (list): List of prefixes
        generations (list): List of generated texts
        model: The language model
        tokenizer: The tokenizer associated with the model
        
    Returns:
        dict: Dictionary containing mean perplexity and individual scores
    """
    perplexities = []

    total = min(len(prefixes), len(generations))
    for prefix, gen in tqdm(zip(prefixes, generations), desc="Calculating perplexity", total=total):
        try:
            ppl = calculate_conditional_perplexity(prefix, gen, model, tokenizer)
            perplexities.append(ppl)
        except Exception as e:
            print(f"Error calculating perplexity: {e}")
            perplexities.append(None)
    
    # Filter out None values for mean calculation
    valid_perplexities = [p for p in perplexities if p is not None]
    mean_ppl = sum(valid_perplexities) / len(valid_perplexities) if valid_perplexities else None
    
    return {
        "mean_perplexity": mean_ppl,
        "individual_perplexities": perplexities
    }

def save_results(results, file_path):
    """
    Save results to a JSON file.
    
    Args:
        results (dict): Dictionary containing perplexity results
        file_path (str): Json file path to save results
    """
    with open(file_path, "w") as f:
        json.dump(results, f, indent=4)