Function init_worker

generate_answers.py:14–27

Initialize worker process, set GPU environment, and initialize the model

(model_path, gpu_id)

from pathlib import Path
import jsonlines
def init_worker(model_path, gpu_id):
    """Initialize worker process, set GPU environment, and initialize the model"""
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    print(f"Process {os.getpid()} using GPU {gpu_id}")

    global model
    model = LLM(
        model=model_path,
        max_num_batched_tokens=8192,
        max_model_len=8192,
        seed=1,
        trust_remote_code=True,
        tensor_parallel_size=1
    )

def process_single_item(item, sampling_params, num_batches):
    """Process a single data item"""