Upload artifact

src/modal_app_rewrite_sota.py

@@ -0,0 +1,238 @@
+"""
+Modal app: SOTA rewriting with BART-SFT-DPO model.
+
+Uses the full trained pipeline (SFT + DPO adversarial) for AI text rewriting.
+Much higher quality than the Qwen2.5-1.5B baseline.
+
+Usage:
+  modal run -q src/modal_app_rewrite_sota.py --text "Your AI text" --verify
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import sys
+import time
+
+import modal
+
+image = (
+    modal.Image.debian_slim(python_version="3.12")
+    .env({"PIP_PROGRESS_BAR": "off", "PYTHONIOENCODING": "utf-8"})
+    .pip_install("torch>=2.4.0", "transformers>=4.45.0", "accelerate>=0.34.0", "huggingface_hub>=0.26.0")
+)
+
+app = modal.App("evasion-detection-sota", image=image)
+hf_cache = modal.Volume.from_name("hf-cache", create_if_missing=True)
+
+MODEL_REPO = "simonlesaumon/evasion-detection-models"
+MODEL_NAME = "bart-sft-style-humanization"  # SFT model (DPO overfit, use SFT directly)
+
+
+@app.function(
+    gpu=os.getenv("MODAL_GPU", "T4"),
+    timeout=60 * 15,
+    scaledown_window=60 * 3,
+    volumes={"/root/.cache/huggingface": hf_cache},
+)
+def rewrite_sota(
+    text: str,
+    verify: bool = True,
+    max_input_length: int = 512,
+    max_output_length: int = 256,
+    temperature: float = 0.8,
+    top_p: float = 0.92,
+    repetition_penalty: float = 1.1,
+) -> dict:
+    """Rewrite AI text using the SOTA BART-SFT model with style embeddings."""
+    import torch
+    import torch.nn as nn
+    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+    print(f"[SOTA] Loading SFT model from {MODEL_REPO} subfolder={MODEL_NAME}...")
+    tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
+
+    # Add style tokens that were used during SFT training
+    style_token_human = "<human>"
+    for tok in [style_token_human]:
+        if tok not in tokenizer.get_vocab():
+            tokenizer.add_tokens([tok])
+
+    # Load the BART base from SFT checkpoint
+    bart = AutoModelForSeq2SeqLM.from_pretrained(
+        MODEL_REPO, subfolder=MODEL_NAME, dtype=torch.float32,
+    )
+    if len(tokenizer) > bart.config.vocab_size:
+        bart.resize_token_embeddings(len(tokenizer))
+
+    # Load style modules
+    from huggingface_hub import hf_hub_download
+    style_path = hf_hub_download(
+        repo_id=MODEL_REPO, filename=f"{MODEL_NAME}/style_modules.pt",
+    )
+    style_ckpt = torch.load(style_path, map_location="cpu")
+
+    # Reconstruct style embeddings
+    hidden_size = bart.config.d_model
+    style_embeddings = nn.Embedding(2, hidden_size)
+    style_proj = nn.Sequential(
+        nn.Linear(hidden_size, 1024), nn.GELU(), nn.Linear(1024, hidden_size),
+    )
+    style_embeddings.load_state_dict(style_ckpt["style_embeddings"])
+    style_proj.load_state_dict(style_ckpt["style_proj"])
+
+    bart.eval()
+    style_embeddings.eval()
+    style_proj.eval()
+
+    if torch.cuda.is_available():
+        bart = bart.to("cuda")
+        style_embeddings = style_embeddings.to("cuda")
+        style_proj = style_proj.to("cuda")
+
+    print(f"[SOTA] Rewriting {len(text.split())} words with human style...")
+    start = time.time()
+
+    # Prepend <human> token to input text
+    input_text = f"{style_token_human} {text}"
+    inputs = tokenizer(
+        input_text, max_length=max_input_length, truncation=True,
+        return_tensors="pt",
+    )
+    if torch.cuda.is_available():
+        inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+    # Get encoder outputs
+    encoder_outputs = bart.model.encoder(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        return_dict=True,
+    )
+
+    # Inject human style embedding
+    human_style_id = torch.tensor([1], device=encoder_outputs.last_hidden_state.device)
+    style_emb = style_embeddings(human_style_id)
+    style_emb = style_proj(style_emb)
+    encoder_outputs.last_hidden_state = (
+        encoder_outputs.last_hidden_state + style_emb.unsqueeze(1)
+    )
+
+    # Generate with style-injected encoder outputs
+    with torch.no_grad():
+        outputs = bart.generate(
+            encoder_outputs=encoder_outputs,
+            attention_mask=inputs["attention_mask"],
+            max_length=max_output_length,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    rewritten = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    elapsed = time.time() - start
+
+    result = {
+        "status": "completed",
+        "model": f"{MODEL_REPO}/{MODEL_NAME}",
+        "original": text,
+        "rewritten": rewritten,
+        "original_words": len(text.split()),
+        "rewritten_words": len(rewritten.split()),
+        "elapsed_seconds": round(elapsed, 2),
+    }
+
+    if verify:
+        result["verification"] = _verify_output(text, rewritten)
+
+    return result
+
+
+def _verify_output(original: str, rewritten: str) -> dict:
+    """Verify rewrite quality."""
+    orig_w = len(original.split())
+    rew_w = len(rewritten.split())
+    ratio = rew_w / max(orig_w, 1)
+    issues, ok = [], []
+
+    if ratio < 0.4:
+        issues.append(f"Too short: {rew_w}w vs {orig_w}w")
+    elif ratio > 2.0:
+        issues.append(f"Too long: {rew_w}w vs {orig_w}w ({ratio:.2f}x)")
+    else:
+        ok.append(f"Length: {orig_w}w -> {rew_w}w ({ratio:.2f}x)")
+
+    artifacts = ["###", "Paraphrase:", "Here is", "Let me know"]
+    found = [a for a in artifacts if a.lower() in rewritten.lower()]
+    if found:
+        issues.append(f"Artifacts: {found}")
+    else:
+        ok.append("No artifacts detected")
+
+    orig_nums = set(re.findall(r'\b\d+\b', original))
+    rew_nums = set(re.findall(r'\b\d+\b', rewritten))
+    missing = orig_nums - rew_nums
+    if missing:
+        issues.append(f"Missing numbers: {missing}")
+    elif orig_nums:
+        ok.append(f"Numbers: {len(orig_nums)}/{len(orig_nums)} preserved")
+
+    return {
+        "passed": len(issues) == 0,
+        "ok": ok, "issues": issues,
+        "length_ratio": round(ratio, 2),
+        "original_words": orig_w, "rewritten_words": rew_w,
+    }
+
+
+@app.local_entrypoint()
+def main(
+    text: str = "",
+    text_file: str = "",
+    gpu: str = "T4",
+    verify: bool = True,
+    output: str = "output/rewrite_sota_result.json",
+):
+    """SOTA rewriting entrypoint."""
+    os.environ["MODAL_GPU"] = gpu
+
+    if text:
+        pass
+    elif text_file:
+        with open(text_file, "r", encoding="utf-8") as f:
+            text = f.read().strip()
+    else:
+        text = "Artificial intelligence has revolutionized natural language processing."
+
+    print("=" * 50)
+    print(f" SOTA Rewrite — BART-DPO")
+    print(f" Model: {MODEL_REPO}/{MODEL_NAME} | GPU: {gpu}")
+    print("=" * 50)
+    print(f"\n[Input] {len(text.split())} words:")
+    print(text[:200] + ("..." if len(text) > 200 else ""))
+
+    result = rewrite_sota.remote(text=text, verify=verify)
+
+    if result.get("status") == "completed":
+        rew = result["rewritten"]
+        print(f"\n--- Rewrite ---")
+        print(f"  Words: {result['original_words']} -> {result['rewritten_words']}")
+        print(f"  Time: {result['elapsed_seconds']}s")
+        print(f"\n  {rew[:500]}")
+
+        if result.get("verification"):
+            v = result["verification"]
+            print(f"\n  Verify: {'OK' if v['passed'] else 'ISSUES'}")
+            for check in v.get("ok", []):
+                print(f"    + {check}")
+            for issue in v.get("issues", []):
+                print(f"    - {issue}")
+
+    os.makedirs(os.path.dirname(output) or ".", exist_ok=True)
+    with open(output, "w", encoding="utf-8") as f:
+        json.dump(result, f, indent=2, ensure_ascii=False, default=str)
+
+    print(f"\n[Save] {output}")