PerturbReason/PerturbReason_dataset_code

1"""2BioMistral Batch Inference with vLLM3Aligned with qwen_batch_vllm pipeline data format.4BioMistral-7B (Mistral architecture) runs efficiently in bf16 on a single A100.5"""6import argparse7import json8import os9import glob10import re11import random12 13 14ANSWER_SUFFIX = 'Answer with either "up", "down", or "unchanged".'15LABEL_ONLY_SUFFIX = ' Output answer label only.'16LABEL_PATTERN = re.compile(r"\b(up|down|unchanged)\b", re.IGNORECASE)17 18 19def parse_args():20    parser = argparse.ArgumentParser(description="BioMistral Batch Inference with vLLM")21    parser.add_argument("--model_path", type=str, required=True,22                        help="Path to the BioMistral model")23    parser.add_argument("--input_dir", type=str, required=True,24                        help="Directory containing input .jsonl files")25    parser.add_argument("--output_dir", type=str, required=True,26                        help="Directory to save output .jsonl files")27    parser.add_argument("--max_new_tokens", type=int, default=1024,28                        help="Max tokens to generate")29    parser.add_argument("--tp_size", type=int, default=1,30                        help="Tensor parallel size (number of GPUs)")31    parser.add_argument("--gpu_memory_utilization", type=float, default=0.9,32                        help="GPU memory utilization for vLLM")33    parser.add_argument("--temperature", type=float, default=0.2,34                        help="Sampling temperature")35    parser.add_argument("--top_p", type=float, default=0.95,36                        help="Top-p sampling parameter")37    parser.add_argument("--repetition_penalty", type=float, default=1.05,38                        help="Penalty to discourage repetitive degeneration")39    parser.add_argument("--limit_files", type=int, default=0,40                        help="If > 0, only process the first N input files")41    parser.add_argument("--limit_records", type=int, default=0,42                        help="If > 0, only process the first N records per input file")43    parser.add_argument("--disable_chat_template", action="store_true",44                        help="Disable BioMistral's instruct chat template wrapper")45    parser.add_argument("--overwrite_existing", action="store_true",46                        help="Overwrite existing output files instead of skipping them")47    parser.add_argument("--guided_decoding", action="store_true", default=True,48                        help="Use vLLM guided decoding to constrain output to the 3 label tokens")49    parser.add_argument("--no_guided_decoding", dest="guided_decoding", action="store_false",50                        help="Disable guided decoding")51    return parser.parse_args()52 53 54def format_generation_prompt(prompt_text, tokenizer, use_chat_template):55    if not use_chat_template:56        return prompt_text57    if "[INST]" in prompt_text:58        return prompt_text59    # Fixed: Added closing parenthesis and proper return60    return tokenizer.apply_chat_template(61        [{"role": "user", "content": prompt_text}],62        tokenize=False,63        add_generation_prompt=True,64    )65 66# This was previously mangled into the function above67LABEL_PATTERN = re.compile(68    r"\b(up(?:regulated?|regulation)?|down(?:regulated?|regulation)?|unchanged)\b",69    re.IGNORECASE,70)71 72 73# Canonical label order — shuffled per call to break positional bias74_LABELS = ["up", "down", "unchanged"]75 76 77def _shuffled_label_str() -> str:78    """Return a randomly ordered label list string, e.g. '"down", "unchanged", or "up"'."""79    order = _LABELS.copy()80    random.shuffle(order)81    quoted = [f'"{l}"' for l in order]82    return ", ".join(quoted[:-1]) + ", or " + quoted[-1]83 84 85def rewrite_prompt_for_label_only(prompt_text):86    # Find where the CoT instruction begins and slice it off cleanly87    marker = "Please reason step by step."88    marker_index = prompt_text.find(marker)89 90    if marker_index != -1:91        base_prompt = prompt_text[:marker_index].strip()92    else:93        # Fallback if the exact marker string isn't found94        base_prompt = prompt_text.strip()95 96    # Randomize label order in the trailing answer line to break positional bias97    # Replace the canonical "up", "down", or "unchanged" order with a shuffled one98    shuffled = _shuffled_label_str()99    for pattern in [100        'Answer with either "up", "down", or "unchanged".',101        'Answer with either "up", "down", or "unchanged"',102        "Answer with either 'up', 'down', or 'unchanged'.",103        "Answer with either 'up', 'down', or 'unchanged'",104    ]:105        if pattern in base_prompt:106            base_prompt = base_prompt.replace(pattern, f"Answer with either {shuffled}.")107            break108 109    # Append a highly restrictive instruction block (also with shuffled order)110    strict_instruction = (111        f"\n\nCRITICAL INSTRUCTION: Output exactly one word. "112        f"Your entire response must be ONLY one of: {shuffled}. "113        "Do not include any explanations, punctuation, or other text."114    )115 116    return base_prompt + strict_instruction117 118 119def normalize_label_output(generated_text):120    cleaned_text = generated_text.strip().lower()121 122    # 1. First check for an exact match (since we instructed it to output one word)123    if cleaned_text in ["up", "down", "unchanged"]:124        return cleaned_text125 126    # 2. Fallback: Check if the string simply starts with the target words127    if cleaned_text.startswith("up"): return "up"128    if cleaned_text.startswith("down"): return "down"129    if cleaned_text.startswith("unchanged"): return "unchanged"130 131    # 3. Last resort: Regex search132    matches = LABEL_PATTERN.findall(cleaned_text)133    if matches:134        last = matches[-1].lower()135        if last.startswith("up"): return "up"136        if last.startswith("down"): return "down"137        return "unchanged"138 139    return cleaned_text140 141 142def main():143    args = parse_args()144    from transformers import AutoTokenizer145    from vllm import LLM, SamplingParams146 147    tokenizer = AutoTokenizer.from_pretrained(148        args.model_path,149        trust_remote_code=True,150    )151    if tokenizer.pad_token is None and tokenizer.eos_token is not None:152        tokenizer.pad_token = tokenizer.eos_token153 154    use_chat_template = bool(getattr(tokenizer, "chat_template", None)) and not args.disable_chat_template155    print(f"Using chat template: {use_chat_template}")156 157    input_files = sorted(glob.glob(os.path.join(args.input_dir, "*.jsonl")))158    if args.limit_files > 0:159        input_files = input_files[:args.limit_files]160 161    # --- 1. Load vLLM ---162    # BioMistral-7B in bf16 needs ~14GB, fits easily on a single A100-40GB.163    # No need for 4-bit quantization — bf16 is faster and more accurate.164    print(f"Loading BioMistral via vLLM from {args.model_path}...")165    llm = LLM(166        model=args.model_path,167        tensor_parallel_size=args.tp_size,168        trust_remote_code=True,169        dtype="bfloat16",170        max_model_len=8192,171        gpu_memory_utilization=args.gpu_memory_utilization,172        enforce_eager=True,173    )174 175    sampling_kwargs = {176        "temperature": args.temperature,177        "top_p": args.top_p,178        "repetition_penalty": args.repetition_penalty,179        "max_tokens": args.max_new_tokens,180    }181    if tokenizer.eos_token_id is not None:182        sampling_kwargs["stop_token_ids"] = [tokenizer.eos_token_id]183 184    # Guided decoding: constrain output to exactly the three label strings.185    # This forces the model to compete among its actual logits for the three186    # choices rather than defaulting to "up" as a common English token.187    if args.guided_decoding:188        try:189            from vllm.sampling_params import GuidedDecodingParams190            sampling_kwargs["guided_decoding"] = GuidedDecodingParams(191                choice=["up", "down", "unchanged"]192            )193            print("Guided decoding enabled via GuidedDecodingParams.")194        except (ImportError, AttributeError):195            try:196                sampling_kwargs["guided_choice"] = ["up", "down", "unchanged"]197                print("Guided decoding enabled via guided_choice (legacy API).")198            except Exception as e:199                print(f"WARNING: Could not enable guided decoding: {e}")200    else:201        print("Guided decoding disabled.")202 203    sampling_params = SamplingParams(**sampling_kwargs)204 205    # --- 2. Process Files ---206    os.makedirs(args.output_dir, exist_ok=True)207    print(f"Found {len(input_files)} JSONL files to process.")208    if args.limit_files > 0:209        print(f"Limiting to first {args.limit_files} input file(s).")210    if args.limit_records > 0:211        print(f"Limiting to first {args.limit_records} record(s) per file.")212 213    for file_path in input_files:214        file_name = os.path.basename(file_path)215        output_path = os.path.join(args.output_dir, f"biomistral_pred_{file_name}")216 217        if os.path.exists(output_path) and not args.overwrite_existing:218            print(f"Skipping {file_name} (output exists).")219            continue220 221        print(f"Processing: {file_name}")222 223        # Load all prompts for this file at once224        prompts = []225        original_records = []226        with open(file_path, 'r', encoding='utf-8') as f:227            for line in f:228                if not line.strip():229                    continue230                try:231                    record = json.loads(line)232                    if record.get("prompt"):233                        prompt_text = rewrite_prompt_for_label_only(record["prompt"])234                        prompts.append(235                            format_generation_prompt(236                                prompt_text,237                                tokenizer,238                                use_chat_template,239                            )240                        )241                        original_records.append(record)242                        if args.limit_records > 0 and len(original_records) >= args.limit_records:243                            break244                except json.JSONDecodeError:245                    continue246 247        if not prompts:248            print(f"  No prompts found in {file_name}; skipping.")249            continue250 251        # --- 3. Batch Inference ---252        # vLLM handles batching internally — much faster than sequential HF generate253        outputs = llm.generate(prompts, sampling_params)254 255        # --- 4. Save Results ---256        with open(output_path, 'w', encoding='utf-8') as f_out:257            for record, output in zip(original_records, outputs):258                generated_text = ""259                if output.outputs:260                    generated_text = normalize_label_output(output.outputs[0].text)261                result_entry = {262                    "source_file": file_name,263                    "prompt": record["prompt"],264                    "ground_truth_response": record.get("response", ""),265                    "model_output": generated_text,266                }267                f_out.write(json.dumps(result_entry) + "\n")268 269        print(f"  Saved {len(original_records)} predictions to {output_path}")270 271    print("\nDone.")272 273 274if __name__ == "__main__":275    main()276