inference.py

#!/usr/bin/env python3
"""DEPO-Paraphrase — standalone inference (no finetune-LLM src/ dependency).

DEPO: Detector-Evasive Paraphrase via Constrained Policy Optimization.
LoRA on Qwen/Qwen3-4B-Instruct-2507 (MAGE detector + BERTScore, checkpoint-750).

Usage:
    python inference.py --adapter_path /path/to/checkpoint-750
    python inference.py --adapter_path YOUR_USERNAME/depo-paraphrase
"""

from __future__ import annotations

import argparse
import os
import re
from dataclasses import dataclass
from typing import Any

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# ---------------------------------------------------------------------------
# Model metadata (matches configs/eval_grpo_rewards.yaml)
# ---------------------------------------------------------------------------

BASE_MODEL_ID = "Qwen/Qwen3-4B-Instruct-2507"

DEFAULT_TEMPLATE = (
    "Paraphrase the following text while preserving its meaning. "
    "Return only the paraphrase.\n\n"
    "Text:\n{source_text}"
)

DEFAULT_MAX_NEW_TOKENS = 512
DEFAULT_TEMPERATURE = 0.9
DEFAULT_TOP_P = 0.95
DEFAULT_TORCH_DTYPE = "bfloat16"

_NOTE_PARENS = re.compile(r"\([^)]*(?:Note|note)\s*:[^)]*\)", re.DOTALL)

_META_LABELS = (
    r"Note|Disclaimer|Important|Reminder|Caution|Warning|"
    r"Editor(?:'|\u2019)s?\s+note|N\.\s*B\.|NB|P\.\s*S\.|PS|"
    r"FYI|Observation|Observations|Clarification|Addendum|Postscript|"
    r"Annotation|Commentary|Correction|Misleading|Inconsistency|"
    r"Contradiction|Meta[\s-]*note"
)
_META_PARAGRAPH_START = re.compile(
    rf"(?:"
    rf"\n{{2,}}\s*(?:#{{1,6}}\s+)?(?:\*\*)?\s*(?:{_META_LABELS})\s*(?:\*\*)?\s*[:\uFF1A.]"
    rf"|\n\s*\*\*\s*(?:{_META_LABELS})\s*\*\*\s*[:\uFF1A.]"
    rf"|\n\s*#{{1,6}}\s+(?:{_META_LABELS})\b\s*[:\uFF1A.]"
    rf")",
    re.IGNORECASE,
)


@dataclass
class ParaphraseModel:
    model: Any
    tokenizer: Any
    template: str = DEFAULT_TEMPLATE

    def format_prompt(self, source_text: str, task_instruction: str = "") -> str:
        body = self.template.format(source_text=source_text)
        if task_instruction.strip():
            return (
                f'The following text was written in response to this task: "{task_instruction}"\n\n'
                + body
            )
        return body

    @staticmethod
    def postprocess(text: str) -> str:
        text = _NOTE_PARENS.sub("", text).strip()
        if not text:
            return text
        last_cut: int | None = None
        for m in _META_PARAGRAPH_START.finditer(text):
            last_cut = m.start()
        if last_cut is not None:
            text = text[:last_cut].rstrip()
        return text.strip()

    def rewrite(
        self,
        source_text: str,
        *,
        task_instruction: str = "",
        max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
        temperature: float = DEFAULT_TEMPERATURE,
        top_p: float = DEFAULT_TOP_P,
    ) -> str:
        prompt = self.format_prompt(source_text, task_instruction=task_instruction)
        messages = [{"role": "user", "content": prompt}]
        chat_input = self.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = self.tokenizer(chat_input, return_tensors="pt").to(self.model.device)
        with torch.no_grad():
            output = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=temperature,
                top_p=top_p,
                pad_token_id=self.tokenizer.pad_token_id,
            )
        prompt_len = inputs["input_ids"].shape[1]
        generated = self.tokenizer.decode(
            output[0][prompt_len:], skip_special_tokens=True
        )
        return self.postprocess(generated)


def _hf_token() -> str | None:
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def load_paraphrase_model(
    adapter_path: str,
    *,
    base_model_id: str = BASE_MODEL_ID,
    torch_dtype: str = DEFAULT_TORCH_DTYPE,
    trust_remote_code: bool = True,
) -> ParaphraseModel:
    dtype = getattr(torch, torch_dtype, torch.bfloat16)
    token = _hf_token()
    tokenizer = AutoTokenizer.from_pretrained(
        adapter_path,
        trust_remote_code=trust_remote_code,
        token=token,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    base = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=dtype,
        trust_remote_code=trust_remote_code,
        device_map="auto",
        token=token,
    )
    model = PeftModel.from_pretrained(
        base, adapter_path, is_trainable=False, token=token
    )
    model.eval()
    return ParaphraseModel(model=model, tokenizer=tokenizer)


# Sample inputs for local smoke test (simulated AI-style passages).
DEMO_SAMPLES = [
    {
        "id": "demo_1",
        "source_text": (
            "Climate change represents one of the most pressing challenges of our time. "
            "Rising global temperatures have led to more frequent extreme weather events, "
            "including hurricanes, droughts, and wildfires. Scientists emphasize that "
            "reducing greenhouse gas emissions is essential to mitigate future impacts."
        ),
    },
    {
        "id": "demo_2",
        "source_text": (
            "Machine learning has transformed numerous industries by enabling systems to "
            "learn patterns from data without explicit programming. Applications range "
            "from medical diagnosis to autonomous vehicles, demonstrating both the "
            "potential and the ethical considerations of automated decision-making."
        ),
    },
]


def main() -> None:
    parser = argparse.ArgumentParser(description="Paraphrase inference (standalone)")
    parser.add_argument(
        "--adapter_path",
        type=str,
        required=True,
        help="Local path or Hugging Face repo id with LoRA adapter weights",
    )
    parser.add_argument(
        "--base_model_id",
        type=str,
        default=BASE_MODEL_ID,
        help=f"Base causal LM (default: {BASE_MODEL_ID})",
    )
    parser.add_argument("--max_new_tokens", type=int, default=DEFAULT_MAX_NEW_TOKENS)
    parser.add_argument("--temperature", type=float, default=DEFAULT_TEMPERATURE)
    parser.add_argument("--top_p", type=float, default=DEFAULT_TOP_P)
    parser.add_argument(
        "--text",
        type=str,
        default=None,
        help="Single text to paraphrase; if omitted, run built-in demo samples",
    )
    args = parser.parse_args()

    print(f"Loading base={args.base_model_id} adapter={args.adapter_path}")
    pm = load_paraphrase_model(
        args.adapter_path,
        base_model_id=args.base_model_id,
    )

    if args.text:
        samples = [{"id": "custom", "source_text": args.text}]
    else:
        samples = DEMO_SAMPLES

    for rec in samples:
        out = pm.rewrite(
            rec["source_text"],
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_p=args.top_p,
        )
        print("\n" + "=" * 72)
        print(f"[{rec['id']}] SOURCE ({len(rec['source_text'])} chars)")
        print(rec["source_text"])
        print(f"\n[{rec['id']}] REWRITE ({len(out)} chars)")
        print(out)


if __name__ == "__main__":
    main()