Function process_data

generate_answers.py:122–240

Process data using vLLM to generate answers. This function provides compatibility with the original pipeline interface. Args: model_path (str): Path to the model input_file (str): Path to input JSON file output_file (str): Path to output JSON file ap

(
    model_path,
    input_file,
    output_file,
    api_port=8012,  # Not used but kept for compatibility
    num_processes=96,  # Not used but kept for compatibility
    batch_size=200,  # This will be used as 'n' (answers per batch)
    save_interval=16,  # Not used but kept for compatibility
    resume=True,  # Will be handled via checkpoint mechanism
    mode=None,  # Not used but kept for compatibility
    num_answers=3200  # This will be used as 'nums_answer' (total answers)
)

    return all_results

def process_data(
    model_path,
    input_file,
    output_file,
    api_port=8012,  # Not used but kept for compatibility
    num_processes=96,  # Not used but kept for compatibility
    batch_size=200,  # This will be used as 'n' (answers per batch)
    save_interval=16,  # Not used but kept for compatibility
    resume=True,  # Will be handled via checkpoint mechanism
    mode=None,  # Not used but kept for compatibility
    num_answers=3200  # This will be used as 'nums_answer' (total answers)
):
    """
    Process data using vLLM to generate answers.
    This function provides compatibility with the original pipeline interface.
    
    Args:
        model_path (str): Path to the model
        input_file (str): Path to input JSON file
        output_file (str): Path to output JSON file
        api_port (int): Not used with vLLM, kept for compatibility
        num_processes (int): Not used with vLLM, kept for compatibility
        batch_size (int): Used as 'n' - number of answers per batch
        save_interval (int): Not used with vLLM, kept for compatibility
        resume (bool): Will use checkpoint mechanism
        mode (str): Not used with vLLM, kept for compatibility
        num_answers (int): Total number of answers to generate per theorem
        
    Returns:
        list: The processed data
    """
    # Setup checkpoint directory
    current_directory = os.getcwd()
    checkpoint_dir = os.path.join(current_directory, 'checkpoint_mp')
    os.makedirs(checkpoint_dir, exist_ok=True)
    
    # Read data
    print(f"Reading data from {input_file}...")
    data = []
    with jsonlines.open(input_file) as reader:
        for obj in reader:
            data.append(obj)
    
    # Calculate num_batches
    n = batch_size  # Use batch_size as 'n'
    nums_answer = num_answers
    num_batches = math.ceil(nums_answer / n)
    
    # Set sampling parameters
    sampling_params = SamplingParams(
        temperature=1.0,
        max_tokens=2048,
        top_p=0.95,
        n=n,
    )
    
    # Get available GPUs
    available_gpus = os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")