PerturbReason/PerturbReason_dataset_code

1"""2NatureLM Batch Inference with vLLM3Aligned with qwen_batch_vllm pipeline data format.4NatureLM is a Mixtral 8x7B MoE model (~90GB in bf16).5Uses 4-bit BitsAndBytes quantization (~23GB) with tp_size=2 across 2x A100-40GB.6vLLM has native MoE support, providing massive speedup over sequential HF generate.7"""8import argparse9import json10import os11import glob12import re13import random14 15# Avoid Torch Inductor's CPU ISA dry-compile probe, which triggered file-lock16# timeouts during previous multi-worker NatureLM startup on the cluster.17os.environ.setdefault("TORCHINDUCTOR_VEC_ISA_OK", "0")18 19# Keep vLLM import LAZY (inside main) so the module-level import does NOT20# trigger CUDA initialisation in the spawned worker processes.  Importing21# vLLM at module level caused `device=1, num_gpus=1` assertion failures in22# multiproc_executor workers on some SLURM allocations (CUDA context was23# cached before the worker's CUDA_VISIBLE_DEVICES was established).24 25LABEL_PATTERN = re.compile(26    r"\b(up(?:regulated?|regulation)?|down(?:regulated?|regulation)?|unchanged)\b",27    re.IGNORECASE,28)29 30_LABELS = ["up", "down", "unchanged"]31 32 33def _shuffled_label_str() -> str:34    """Return a randomly ordered label list string, e.g. '"down", "unchanged", or "up"'."""35    order = _LABELS.copy()36    random.shuffle(order)37    quoted = [f'"{l}"' for l in order]38    return ", ".join(quoted[:-1]) + ", or " + quoted[-1]39 40 41def rewrite_prompt_for_label_only(prompt_text):42    # Strip the CoT / structured-response instruction (same marker as BioMistral)43    marker = "Please reason step by step."44    marker_index = prompt_text.find(marker)45    if marker_index != -1:46        base_prompt = prompt_text[:marker_index].strip()47    else:48        base_prompt = prompt_text.strip()49 50    # Randomize label order in the trailing answer line to break positional bias51    shuffled = _shuffled_label_str()52    for pattern in [53        'Answer with either "up", "down", or "unchanged".',54        'Answer with either "up", "down", or "unchanged"',55        "Answer with either 'up', 'down', or 'unchanged'.",56        "Answer with either 'up', 'down', or 'unchanged'",57    ]:58        if pattern in base_prompt:59            base_prompt = base_prompt.replace(pattern, f"Answer with either {shuffled}.")60            break61 62    strict_instruction = (63        f"\n\nCRITICAL INSTRUCTION: Output exactly one word. "64        f"Your entire response must be ONLY one of: {shuffled}. "65        "Do not include any explanations, punctuation, or other text."66    )67    return base_prompt + strict_instruction68 69 70def normalize_label_output(generated_text):71    cleaned_text = generated_text.strip()72    if not cleaned_text:73        return cleaned_text74    matches = LABEL_PATTERN.findall(cleaned_text)75    if not matches:76        return cleaned_text77    last = matches[-1].lower()78    if last.startswith("up"):79        return "up"80    if last.startswith("down"):81        return "down"82    return "unchanged"83 84 85def parse_args():86    parser = argparse.ArgumentParser(description="NatureLM Batch Inference with vLLM")87    parser.add_argument("--model_path", type=str, required=True,88                        help="Path to the NatureLM model")89    parser.add_argument("--input_dir", type=str, required=True,90                        help="Directory containing input .jsonl files")91    parser.add_argument("--output_dir", type=str, required=True,92                        help="Directory to save output .jsonl files")93    parser.add_argument("--max_new_tokens", type=int, default=5,94                        help="Max tokens to generate")95    parser.add_argument("--tp_size", type=int, default=2,96                        help="Tensor parallel size (default 2 for MoE on 2x A100-40GB)")97    parser.add_argument("--gpu_memory_utilization", type=float, default=0.9,98                        help="GPU memory utilization for vLLM")99    parser.add_argument("--temperature", type=float, default=0.1,100                        help="Sampling temperature")101    parser.add_argument("--top_p", type=float, default=1.0,102                        help="Top-p sampling parameter")103    parser.add_argument("--limit_files", type=int, default=0,104                        help="If > 0, only process the first N input files")105    parser.add_argument("--limit_records", type=int, default=0,106                        help="If > 0, only process the first N records per input file")107    parser.add_argument("--overwrite_existing", action="store_true",108                        help="Overwrite existing output files instead of skipping them")109    parser.add_argument("--guided_decoding", action="store_true", default=True,110                        help="Constrain output to the 3 label tokens via vLLM guided decoding")111    parser.add_argument("--no_guided_decoding", dest="guided_decoding", action="store_false",112                        help="Disable guided decoding")113    return parser.parse_args()114 115 116def main():117    args = parse_args()118    # Lazy import: keeps vLLM out of the module-level namespace so spawned119    # workers do not trigger CUDA initialisation before their device is set.120    from vllm import LLM, SamplingParams121 122    if not os.path.isdir(args.input_dir):123        raise FileNotFoundError(f"Input directory does not exist: {args.input_dir}")124 125    input_files = sorted(glob.glob(os.path.join(args.input_dir, "*.jsonl")))126    if args.limit_files > 0:127        input_files = input_files[:args.limit_files]128 129    if not input_files:130        print(f"ERROR: No .jsonl files found in {args.input_dir}")131        return132 133    # --- 1. Load vLLM ---134    # NatureLM is Mixtral 8x7B (~93GB bf16). With 4-bit BitsAndBytes quantization135    # it's ~23GB total. We use tp_size=2 across 2x A100-40GB (~12GB/GPU) because136    # MoE layers dequantize weights on-the-fly during forward pass, causing memory137    # spikes that OOM on a single 40GB GPU.138    print(f"Loading NatureLM (Mixtral 8x7B, 4-bit) via vLLM from {args.model_path}...")139    llm = LLM(140        model=args.model_path,141        tensor_parallel_size=args.tp_size,142        trust_remote_code=True,143        quantization="bitsandbytes",144        load_format="bitsandbytes",145        max_model_len=8192,146        gpu_memory_utilization=args.gpu_memory_utilization,147        enforce_eager=True,148    )149 150    sampling_kwargs = {151        "temperature": args.temperature,152        "top_p": args.top_p,153        "max_tokens": args.max_new_tokens,154    }155 156    # Guided decoding: constrain output to exactly the three label strings.157    # NatureLM (Mixtral MoE) tends to generate long CoT; this forces158    # genuine per-label logit competition after the prompt.159    if args.guided_decoding:160        try:161            from vllm.sampling_params import GuidedDecodingParams162            sampling_kwargs["guided_decoding"] = GuidedDecodingParams(163                choice=["up", "down", "unchanged"]164            )165            print("Guided decoding enabled via GuidedDecodingParams.")166        except (ImportError, AttributeError):167            try:168                sampling_kwargs["guided_choice"] = ["up", "down", "unchanged"]169                print("Guided decoding enabled via guided_choice (legacy API).")170            except Exception as e:171                print(f"WARNING: Could not enable guided decoding: {e}")172    else:173        print("Guided decoding disabled.")174 175    sampling_params = SamplingParams(**sampling_kwargs)176 177    # --- 2. Process Files ---178    os.makedirs(args.output_dir, exist_ok=True)179    print(f"Found {len(input_files)} JSONL files to process.")180    if args.limit_files > 0:181        print(f"Limiting to first {args.limit_files} input file(s).")182    if args.limit_records > 0:183        print(f"Limiting to first {args.limit_records} record(s) per file.")184 185    for file_path in input_files:186        file_name = os.path.basename(file_path)187        output_path = os.path.join(args.output_dir, f"naturelm_pred_{file_name}")188 189        if os.path.exists(output_path) and not args.overwrite_existing:190            print(f"Skipping {file_name} (output exists).")191            continue192 193        print(f"Processing: {file_name}")194 195        # Load all prompts for this file at once196        prompts = []197        original_records = []198        with open(file_path, 'r', encoding='utf-8') as f:199            for line in f:200                if not line.strip():201                    continue202                try:203                    record = json.loads(line)204                    if record.get("prompt"):205                        prompt_text = rewrite_prompt_for_label_only(record["prompt"])206                        prompts.append(prompt_text)207                        original_records.append(record)208                        if args.limit_records > 0 and len(original_records) >= args.limit_records:209                            break210                except json.JSONDecodeError:211                    continue212 213        if not prompts:214            print(f"  No prompts found in {file_name}; skipping.")215            continue216 217        # --- 3. Batch Inference ---218        # vLLM handles batching internally — orders of magnitude faster than219        # sequential HF generate() for 8x7B MoE220        outputs = llm.generate(prompts, sampling_params)221 222        # --- 4. Save Results ---223        with open(output_path, 'w', encoding='utf-8') as f_out:224            for record, output in zip(original_records, outputs):225                generated_text = ""226                if output.outputs:227                    generated_text = normalize_label_output(output.outputs[0].text)228                result_entry = {229                    "source_file": file_name,230                    "prompt": record["prompt"],231                    "ground_truth_response": record.get("response", ""),232                    "model_output": generated_text,233                }234                f_out.write(json.dumps(result_entry) + "\n")235 236        print(f"  Saved {len(original_records)} predictions to {output_path}")237 238    print("\nDone.")239 240 241if __name__ == "__main__":242    main()243