src/modal_app_rewrite.py

"""
Modal app: Sentence-by-sentence AI text rewriting for detector evasion.

Instead of rewriting the entire text at once (unreliable on small models),
we split into sentences, paraphrase each one independently, and reassemble.
This prevents hallucinations, maintains length, and keeps facts intact.

Key anti-hallucination features:
- Dynamic max_tokens per sentence (scaled to input length)
- Repetition penalty to prevent rambling
- Per-sentence hallucination detection with fallback to original
- Strict length ratio check (must be < 1.4x)
- Number and proper noun preservation checks

Uses Qwen2.5-1.5B-Instruct on T4 (~0.60/h).

Usage:
  modal run -q src/modal_app_rewrite.py --text "Your AI text"
  modal run -q src/modal_app_rewrite.py --text-file data/bitcoin_text.txt --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")
)

app = modal.App("evasion-detection-rewrite", image=image)


# ---------------------------------------------------------------------------
# Core rewriting — sentence by sentence
# ---------------------------------------------------------------------------

PARAPHRASE_PROMPT = (
    "{sentence}\nIn other words,"
)


def _split_sentences(text: str) -> list[str]:
    """Split text into sentences, preserving paragraphs."""
    raw = re.split(r'(?<=[.!?])\s+', text)
    return [s.strip() for s in raw if s.strip() and len(s.split()) >= 3]


def _tokens_for_sentence(sentence: str) -> int:
    """Dynamic max tokens: scale to input sentence length.

    Short sentence (5w) -> 20 tokens. Long sentence (25w) -> 90 tokens.
    Higher ceiling = more natural sentence length variation (burstiness).
    """
    words = len(sentence.split())
    return max(20, min(int(words * 3.5), 90))


def _extract_key_tokens(text: str) -> set[str]:
    """Extract numbers, proper nouns, and domain terms that must be preserved."""
    tokens = set()
    # Numbers (years, amounts, counts)
    for m in re.finditer(r'\b\d+(?:[.,]\d+)?(?:\s*(?:million|billion|%|years?|USD|BTC|EUR))?\b', text, re.IGNORECASE):
        tokens.add(m.group().strip().lower())
    # Proper nouns (capitalized words, 3+ chars)
    for m in re.finditer(r'\b[A-Z][a-z]{2,}(?:\s+[A-Z][a-z]{2,})*\b', text):
        tokens.add(m.group().strip().lower())
    # Key domain terms (acronyms like BTC, AI, etc.)
    for m in re.finditer(r'\b[A-Z]{2,6}\b', text):
        tokens.add(m.group().strip().lower())
    return tokens


def _rewrite_text(
    text: str,
    model_name: str = "Qwen/Qwen2.5-1.5B-Instruct",
    temperature: float = 0.85,
    top_p: float = 0.94,
    max_new_tokens: int = 60,
    repetition_penalty: float = 1.12,
    dry_run: bool = False,
) -> dict:
    """Rewrite text sentence by sentence."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if dry_run:
        return {"status": "dry_run", "model": model_name}

    print(f"[Rewrite] Loading {model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()

    sentences = _split_sentences(text)
    print(f"[Rewrite] Split into {len(sentences)} sentences")

    paraphrased = []
    total_tokens = 0
    fallback_count = 0
    start = time.time()

    for i, sent in enumerate(sentences):
        prompt = PARAPHRASE_PROMPT.replace("{sentence}", sent)
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        # Dynamic token budget per sentence
        sent_tokens = _tokens_for_sentence(sent)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=sent_tokens,
                temperature=temperature,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        input_len = inputs["input_ids"].shape[1]
        generated = outputs[0][input_len:]
        para = tokenizer.decode(generated, skip_special_tokens=True)

        # Clean + verify fidelity per sentence
        para = _clean_sentence(para, sent)
        if _is_hallucination(para, sent, full_original=text):
            para = sent  # fallback to original
            fallback_count += 1

        paraphrased.append(para)
        total_tokens += len(generated)

        if (i + 1) % 5 == 0:
            print(f"[Rewrite] {i+1}/{len(sentences)} sentences done")

    elapsed = time.time() - start
    rewritten = " ".join(paraphrased)

    print(f"[Rewrite] Done: {len(text.split())}w -> {len(rewritten.split())}w, "
          f"{total_tokens} tokens, {elapsed:.1f}s")

    return {
        "status": "completed",
        "model": model_name,
        "original": text,
        "rewritten": rewritten,
        "original_words": len(text.split()),
        "rewritten_words": len(rewritten.split()),
        "num_sentences": len(sentences),
        "tokens": total_tokens,
        "elapsed_seconds": round(elapsed, 1),
        "fallback_count": fallback_count,
    }


def _clean_full_text(rewritten: str, original: str) -> str:
    """Clean full-text rewrite output."""
    text = rewritten.strip()

    # Remove prompt echoes
    cut_patterns = [
        r"\n\s*TEXT:", r"\n\s*HUMAN REWRITE:",
        r"Here is", r"I hope", r"Let me know", r"Feel free",
        r"Note that", r"Please note", r"What are your thoughts",
    ]
    for pat in cut_patterns:
        m = re.search(pat, text)
        if m:
            text = text[: m.start()].strip()

    # Trim trailing incomplete sentence
    if text and text[-1] not in '.!?"' "'":
        last_end = -1
        for punct in ['. ', '! ', '? ']:
            idx = text.rfind(punct)
            if idx > last_end:
                last_end = idx
        if last_end > len(text) * 0.6:
            text = text[: last_end + 1].strip()

    # Remove trailing incomplete word
    if text and text[-1] not in '.!?"' "'" and ' ' in text:
        last_space = text.rfind(' ')
        if last_space > len(text) * 0.8:
            text = text[:last_space].strip()

    if len(text.split()) < 5:
        return original  # too short, use original

    return text


def _clean_sentence(paraphrased: str, original: str) -> str:
    """Clean a paraphrased sentence — keep ONLY the first sentence, strip artifacts."""
    text = paraphrased.strip()

    # Capitalize first letter (completion format often produces lowercase start)
    if text and text[0].islower():
        text = text[0].upper() + text[1:]

    # STRICT truncation: keep only the first complete sentence
    # Use the FIRST sentence boundary found (not the last)
    first_cut = len(text)
    for punct in ['. ', '! ', '? ']:
        idx = text.find(punct)
        if 5 < idx < first_cut:
            first_cut = idx + 1

    if first_cut < len(text):
        text = text[:first_cut].strip()
    else:
        # No sentence boundary found — trim trailing incomplete word
        # (model ran out of tokens mid-word, e.g. "you mi" → "you")
        if text and not text[-1] in '.!?':
            # Find last space and cut there (keep complete words only)
            last_space = text.rfind(' ')
            if last_space > len(text) * 0.5:  # only if we keep most of the text
                text = text[:last_space].strip()

    # Remove common artifacts (model meta-text)
    cut_patterns = [
        r"\n", r"Paraphrased:", r"Sentence:", r"Output:",
        r"Here is", r"I hope", r"Let me know", r"Feel free",
        r"Note that", r"Please note", r"Keep in mind",
    ]
    for pat in cut_patterns:
        m = re.search(pat, text)
        if m:
            text = text[: m.start()].strip()

    # If empty or just punctuation, fallback
    if len(text.split()) < 2:
        return original

    return text


def _is_hallucination(paraphrased: str, original_sentence: str, full_original: str = "") -> bool:
    """Detect if paraphrase is clearly broken (NOT if it just uses different words).

    Only flag OBJECTIVE failures — a paraphrase IS allowed to use different
    vocabulary. This is the whole point of evasion rewriting.

    Returns True only for clear malfunctions (empty output, meta-text, etc.).
    """
    para_w = len(paraphrased.split())
    orig_w = len(original_sentence.split())

    # 1. Length EXPLOSION: >4x original (real hallucination, not just verbose)
    if orig_w >= 5 and para_w > orig_w * 4:
        return True

    # 2. Numbers lost: if original has specific numbers, they must be preserved
    orig_nums = set(re.findall(r'\b\d+\b', original_sentence))
    para_nums = set(re.findall(r'\b\d+\b', paraphrased))
    if orig_nums and not (orig_nums & para_nums):
        return True

    # 3. Starts with lowercase = model continued prompt, didn't generate
    if paraphrased and paraphrased[0].islower():
        return True

    # 4. Empty/gave up
    if para_w < 2:
        return True

    # 5. Meta-text artifacts (model outputting instructions, not content)
    meta_markers = [
        "here is", "i hope", "let me know", "feel free", "note that",
        "keep in mind", "please note", "i'd be happy", "can i help",
        "would you like", "here's a", "here are some",
        "what are your thoughts", "what do you think", "in conclusion",
        "to sum up", "in summary", "to conclude", "as a final",
        "do you agree", "what's your opinion", "let's discuss",
    ]
    # Also flag if the paraphrase is a QUESTION (ends with ?)
    if paraphrased.strip().endswith('?'):
        return True
    para_lower = paraphrased.lower()
    for marker in meta_markers:
        if marker in para_lower:
            return True

    # 6. Repetition loop: same 2-3 word phrase repeated 4+ times
    words = paraphrased.split()
    if len(words) >= 8:
        for n in [2, 3]:
            for i in range(len(words) - n * 3):
                phrase = " ".join(words[i:i+n])
                count = sum(1 for j in range(len(words) - n + 1)
                           if " ".join(words[j:j+n]) == phrase)
                if count >= 4:
                    return True

    # 7. Too short: paraphrase lost most of the original information
    if orig_w >= 8 and para_w < orig_w * 0.3:
        return True  # e.g., 16-word sentence → "It's decentralized." (2 words)

    # 8. Complete topic shift: zero words of 3+ chars overlap.
    #    Only catches extreme cases (e.g., "wallet" instead of "blockchain").
    if orig_w >= 6 and para_w >= 4:
        orig_words_set = {w.strip(",.!?;:'\"()[]").lower()
                         for w in original_sentence.split()
                         if len(w.strip(",.!?;:'\"()[]")) >= 3}
        para_words_set = {w.strip(",.!?;:'\"()[]").lower()
                         for w in paraphrased.split()
                         if len(w.strip(",.!?;:'\"()[]")) >= 3}
        if orig_words_set and para_words_set:
            if not (orig_words_set & para_words_set):
                return True  # completely different vocabulary = topic shift

    return False


# ---------------------------------------------------------------------------
# Verification
# ---------------------------------------------------------------------------

def _verify_output(original: str, rewritten: str) -> dict:
    """Verify rewrite quality — strict checks for faithfulness."""
    orig_w = len(original.split())
    rew_w = len(rewritten.split())
    ratio = rew_w / max(orig_w, 1)
    issues, ok = [], []

    # 1. Length — strict: flag above 1.4x (hallucination risk)
    if ratio < 0.4:
        issues.append(f"Too short: {rew_w}w vs {orig_w}w")
    elif ratio > 1.4:
        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)")

    # 2. Artifacts
    artifacts = ["###", "Paraphrase:", "Here is", "Let me know", "I hope",
                 "Feel free", "Do you have", "Output:", "Sentence:",
                 "Note that", "Keep in mind", "It's worth noting"]
    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")

    # 3. Numbers preserved — ANY number loss is an issue
    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")
    else:
        ok.append("No numbers to preserve")

    # 4. Key terms preserved (proper nouns ONLY, not sentence-start words)
    # Exclude common English words accidentally capitalized at sentence start
    _COMMON_STARTERS = {
        "this", "that", "these", "those", "instead", "because", "despite",
        "unlike", "some", "many", "while", "although", "however", "therefore",
        "thus", "hence", "nonetheless", "nevertheless", "otherwise", "rather",
        "indeed", "furthermore", "moreover", "besides", "according", "during",
        "after", "before", "since", "until", "when", "where", "whereas",
        "one", "two", "three", "first", "second", "third", "another",
    }
    orig_proper = {p for p in re.findall(r'\b[A-Z][a-z]{2,}\b', original)
                   if p.lower() not in _COMMON_STARTERS}
    rew_lower = rewritten.lower()
    missing_proper = [p for p in orig_proper if p.lower() not in rew_lower]
    if len(missing_proper) > max(1, len(orig_proper) * 0.3):
        issues.append(f"Missing terms: {missing_proper}")
    elif orig_proper:
        ok.append(f"Terms: {len(orig_proper) - len(missing_proper)}/{len(orig_proper)} preserved")
    else:
        ok.append("No proper nouns to track")

    # 5. Hallucinated content: new proper nouns not in original
    rew_proper = {p for p in re.findall(r'\b[A-Z][a-z]{2,}\b', rewritten)
                  if p.lower() not in _COMMON_STARTERS}
    orig_lower = original.lower()
    new_terms = [p for p in rew_proper if p.lower() not in orig_lower]
    if len(new_terms) > 3:
        issues.append(f"Hallucinated terms: {new_terms}")

    # 6. Hallucinated dollar amounts
    if '$' in rewritten and '$' not in original:
        issues.append("Hallucinated dollar amounts")

    return {
        "passed": len(issues) == 0,
        "ok": ok, "issues": issues,
        "length_ratio": round(ratio, 2),
        "original_words": orig_w, "rewritten_words": rew_w,
    }


# ---------------------------------------------------------------------------
# Modal function
# ---------------------------------------------------------------------------

@app.function(
    gpu=os.getenv("MODAL_GPU", "T4"),
    timeout=60 * 15,
    scaledown_window=60 * 3,
)
def rewrite_on_gpu(
    texts: list[str],
    model_name: str = "Qwen/Qwen2.5-1.5B-Instruct",
    temperature: float = 0.85,
    top_p: float = 0.94,
    max_new_tokens: int = 60,
    repetition_penalty: float = 1.12,
    verify: bool = False,
    dry_run: bool = False,
) -> dict:
    """Rewrite texts on Modal GPU."""
    if dry_run:
        return {"status": "dry_run", "num_samples": len(texts)}

    results = []
    for text in texts:
        result = _rewrite_text(
            text=text, model_name=model_name,
            temperature=temperature, top_p=top_p,
            max_new_tokens=max_new_tokens,
            repetition_penalty=repetition_penalty,
            dry_run=False,
        )
        if verify and result.get("status") == "completed":
            result["verification"] = _verify_output(text, result["rewritten"])
        results.append(result)

    return {
        "status": "completed",
        "num_samples": len(results),
        "config": {
            "model": model_name, "temperature": temperature,
            "top_p": top_p, "max_new_tokens": max_new_tokens,
            "repetition_penalty": repetition_penalty,
        },
        "results": results,
    }


# ---------------------------------------------------------------------------
# Local entrypoint
# ---------------------------------------------------------------------------

@app.local_entrypoint()
def main(
    text: str = "",
    text_file: str = "",
    model: str = "Qwen/Qwen2.5-1.5B-Instruct",
    gpu: str = "T4",
    temperature: float = 0.85,
    top_p: float = 0.94,
    max_tokens: int = 60,
    repetition_penalty: float = 1.12,
    verify: bool = True,
    dry_run: bool = False,
    output: str = "output/rewrite_result.json",
):
    """Rewrite AI text sentence by sentence."""
    os.environ["MODAL_GPU"] = gpu

    if text:
        texts = [text]
    elif text_file:
        with open(text_file, "r", encoding="utf-8") as f:
            texts = [f.read().strip()]
    else:
        texts = ["Artificial intelligence has revolutionized natural language processing."]

    print("=" * 50)
    print(f" AI Text Rewrite — Sentence-by-sentence")
    print(f" Model: {model} | GPU: {gpu}")
    print(f" Temp: {temperature} | Top-p: {top_p} | Rep-pen: {repetition_penalty}")
    print("=" * 50)

    for i, t in enumerate(texts):
        print(f"\n[Input {i+1}] {len(t.split())} words:")
        print(t[:200] + ("..." if len(t) > 200 else ""))

    if dry_run:
        result = {"status": "dry_run_validated"}
    else:
        result = rewrite_on_gpu.remote(
            texts=texts, model_name=model,
            temperature=temperature, top_p=top_p,
            max_new_tokens=max_tokens,
            repetition_penalty=repetition_penalty,
            verify=verify, dry_run=False,
        )

    if result.get("status") == "completed":
        for i, r in enumerate(result["results"]):
            rew = r["rewritten"]
            print(f"\n--- Sample {i+1} ---")
            print(f"  Words: {r['original_words']} -> {r['rewritten_words']}")
            print(f"  Sentences: {r.get('num_sentences', '?')}")
            print(f"  Time: {r['elapsed_seconds']}s")
            print(f"\n  {rew[:500]}")

            if r.get("verification"):
                v = r["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}")