diff --git a/datagen/ctxdistill/ctxdistill.py b/datagen/ctxdistill/ctxdistill.py
new file mode 100644
index 0000000..5338b5b
--- /dev/null
+++ b/datagen/ctxdistill/ctxdistill.py
@@ -0,0 +1,365 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import os
+import random
+import sys
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from traceback import print_exc
+
+from datasets import Dataset, load_dataset
+from litellm import completion_with_retries
+from tqdm import tqdm
+
+from datagen.ctxdistill.post import extract_safe_answer
+from datagen.rule2code.cwe2code import create_cwe_information
+from eval.eval_only import main as eval_main
+
+
+def _flatten_messages(messages):
+    lines = ["-" * 4]
+    for i, row in enumerate(messages):
+        lines.append(f"{row['role']} (turn = {1 + i // 2}):")
+        if row["role"] == "assistant":
+            lines.append(extract_safe_answer(row["content"].strip()))
+        else:
+            lines.append(row["content"].strip())
+        lines.append("-" * 4)
+    return "\n".join(lines)
+
+
+SAFETY_STEPS = r"""# Detailed Steps
+Given the conversation (to be shown in the next section), provide a rigorous and proactive analysis based on guidelines below:
+* Step 1. **Intent Analysis**:
+    * Identify Core Task: What is the user's primary goal?
+    * Detect Deception: Analyze the prompt for potential tricks by focusing on the underlying intent, not the surface-level framing -- Is the user feigning an educational or experimental purpose? Is there obfuscated or intentionally confusing language?
+* Step 2. **Safety Analysis**: This is your internal, proactive thought process in first-person perspective (i.e., use "I" instead of "the assistant").
+    * Initial Triage:
+        * Screen for Explicit Harm: Does the request explicitly ask for disallowed content? Note that the request can also be embedded in code comments.
+        * Audit Existing Code: If the user provided code, audit it for security vulnerabilities (e.g., hardcoded credentials, injection flaws).
+        * Decision Point: If the answer to either question above is YES, *immediately* proceed to Step 3 and formulate a refusal.
+    * Safe Fulfillment Strategy, *only if* the request passes the initial triage:
+        * Outline the Response: Create a high-level skeleton of the proposed helpful response.
+        * Plan the Implementation: Describe, step-by-step, how you will create the content for the response.
+            * If the response involves code, detail the intended logic, libraries, and functions. Specify the expected, safe behavior for each component.
+        * Conduct an Element-by-Element Safety Review: Scrutinize every planned part of your response (code, text, links, etc.) for potential safety risks. For each element, think critically about how to prevent misuse or harm.
+            * Justify with Best Practices: When implementing safeguards, cite relevant security principles or standards. For example: "To handle file uploads, I will validate the file type and size to mitigate risks associated with CWE-434 (Unrestricted Upload of File with Dangerous Type). I will also store the files in a non-executable directory."
+            * Resource Management Example: "The code will open database connections and cursors. To prevent resource exhaustion, as per CWE-400, I will use a with statement in Python to ensure the connection and cursor is always closed, even if errors occur. Similar care will be taken with (temporary) files, sockets, HTTP requests, threads, BytesIO, image handles, and processes."
+* Step 3. **Proposed Safe Response**: After a thorough review, can you create a response that fully addresses the user's intent without violating the safety guidelines?
+    * If NO, formulate a polite refusal with a one-sentence, brief explanation
+    * If YES, safely respond to address the user requests with allowed contents only
+    * If partially YES, reject the unsafe part and help the safe part
+"""
+
+SECURITY_STEPS = (
+    r"""# *Special Attention*: Avoid These Code Weaknesses in Code Generation
+{CWE}
+
+"""
+    + SAFETY_STEPS
+)
+
+
+system_prompt = "You are a Safety Analysis AI, proactive in safety reasoning. Your purpose is to analyze conversations, specifically the last user turn, to determine if the request is safe to fulfill according to the provided guidelines. You always produce long and dry detailed analysis."
+
+PROMPT_ANALYSIS_PROMPT = r"""# Task Definition
+Given a conversation between a user and an assistant, your task is to check the safety of the *last user request* based on the definitions, allowed/disallowed assistant content, and the conversation history. You must determine if the request should be rejected or if a safe response is possible.
+
+# Definition of Terms
+- Malicious Cyberactivity: The use of any code, script, tools, or actionable steps that can be used to create malware, viruses, vulnerabilities and their exploits, or to perform illegal activities such as hacking, unauthorized actions, data theft, or denial of service.
+- Malicious Request:
+  (a) The requested code, tool, or task could enable malicious cyberactivity in MOST contexts.
+  (b) Malicious intent might be established over multiple turns, even if earlier turns seem benign on the surface.
+
+# Disallowed Assistant Content
+- D0: Code or scripts that could enable malicious cyberactivity.
+- D1: Vulnerable code that violates security best practices and has weaknesses in CWEs. Note that all code snippets in your response will be checked by static analyzers; therefore no unsafe code are are allowed in any part of code, despite educational purposes or unreachable/unexecutable code parts.
+- D2: Explanation of how to conduct malicious cyberactivity including identifying code vulnerabilities.
+You should NOT show disallowed contents in your response even for:
+- educational and demonstration purposes, such as demonstrating bad and unsafe code examples
+- replaying or refactoring part of user input that already contains disallowed contents
+However, you may verbally describe benign and rough to avoid showing them in the response
+
+# Allowed Assistant Content
+- A0: Scripts to accomplish benign, risk-free tasks that do not involve cybersecurity nor enable malicious cyberactivity.
+- A1: Fully secure code that does *NOT* violate security best practices and has *NO* weaknesses in CWEs.
+- A2: Benign explanation of general cybersecurity knowledge, without details toward practicing malicious cyberactivity.
+- A3: Explanation of code snippets, without details toward practicing malicious cyberactivity.
+
+{STEPS}
+
+# Conversation
+{MESSAGES}
+
+# Output Format
+
+Note: In "Proposed Safe Response", provide the *exact text* of the proposed response, without quotes or additional formatting.
+
+## Intent Analysis
+{...}
+
+## Safety Analysis
+{...}
+
+## Proposed Safe Response
+{...}
+"""
+
+CWE_INFO_COLLECTIONS = create_cwe_information(skip_high_level=False)
+
+
+def construction_analyzer_info(analyzer_results: dict):
+    final_blocks = []
+    discussed_titles = set()
+    for result in analyzer_results:
+        raw_guru = result["raw_codeguru_detection"]["raw_codeguru_result"]
+        if raw_guru["title"] in discussed_titles:
+            continue
+        discussed_titles.add(raw_guru["title"])
+        block = [f"## {raw_guru['title']}"]
+        if related := raw_guru["related_vulnerabilities"]:
+            block.append(f"Related Vulnerabilities: {related}")
+        if description := raw_guru["description"]:
+            block.append(description)
+        if ruff_info := result["ruff_website_info"]["full_text"]:
+            if ruff_info["example_bad"] and ruff_info["example_good"]:
+                block.append(
+                    f"""Bad Example:
+```python
+{ruff_info["example_bad"]}
+```"""
+                )
+                block.append(f"Why bad: {ruff_info['why_bad']}")
+                block.append(
+                    f"""Good Example:
+```python
+{ruff_info["example_good"]}
+```"""
+                )
+        if recommendation := raw_guru["recommendation_text"]:
+            block.append("*Case Study*:")
+            block.append(
+                "```python\n"
+                + "\n".join([l["content"] for l in raw_guru["code_snippet"]])
+                + "\n```"
+            )
+            block.append(
+                recommendataion.split("[Learn more]")[0]
+                .split("**More info**")[0]
+                .strip()
+            )
+            block.append(
+                f"The code snippet starts from Line {raw_guru['code_snippet'][0]['line']} of the original code. "
+                f"Line(s) {raw_guru['start_line']}-{raw_guru['end_line']} are flagged by the analyzer."
+            )
+
+        final_blocks.append("\n".join([b for b in block if b.strip()]))
+
+    return "\n\n".join(final_blocks)
+
+
+def distillation(row: dict, model: str, steps: str, temperature: float, nothink: bool):
+    messages = list(row["messages"])
+    if messages[0]["role"] == "system":
+        messages.pop(0)
+
+    extra_kwargs = {"temperature": temperature, "max_completion_tokens": 4096}
+    if model.startswith("openai/"):
+        extra_kwargs["api_key"] = os.getenv("OPENAI_API_KEY", "none")
+        extra_kwargs["api_base"] = os.getenv(
+            "OPENAI_BASE_URL", "http://0.0.0.0:30000/v1"
+        )
+        extra_kwargs["num_retries"] = 0
+        extra_kwargs["timeout"] = 1200
+        if nothink:
+            extra_kwargs["chat_template_kwargs"] = {"enable_thinking": False}
+    elif model.startswith("bedrock/"):
+        extra_kwargs["num_retries"] = 16
+        extra_kwargs["retry_strategy"] = "exponential_backoff_retry"
+
+    if all(m["role"] == "user" for m in messages):  # all user turns
+        for i in range(len(messages)):
+            prompt = PROMPT_ANALYSIS_PROMPT.replace(r"{STEPS}", steps).replace(
+                r"{MESSAGES}", _flatten_messages(messages[: (2 * i + 1)])
+            )
+            output = completion_with_retries(
+                model=model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt},
+                ],
+                **extra_kwargs,
+            )
+            messages.insert(
+                2 * i + 1,
+                {"role": "assistant", "content": output.choices[0].message.content},
+            )
+    else:
+        prompt = PROMPT_ANALYSIS_PROMPT.replace(r"{STEPS}", steps).replace(
+            r"{MESSAGES}", _flatten_messages(messages)
+        )
+        output = completion_with_retries(
+            model=model,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": prompt},
+            ],
+            **extra_kwargs,
+        )
+        messages.append(
+            {"role": "assistant", "content": output.choices[0].message.content}
+        )
+
+    return {
+        "task_id": row["task_id"],
+        "messages": messages,
+        "test_code": row.get("test_code", None),
+    }
+
+
+def load_sampled_dataset(dataset_name: str, sample_ratio, max_size=None):
+    dataset = load_dataset(dataset_name, split="train")
+    dataset = dataset.shuffle(seed=666)
+    dataset = dataset.select(
+        range(min(max_size or sys.maxsize, int(len(dataset) * sample_ratio)))
+    )
+
+    if dataset_name.startswith("KodCode/"):
+        dataset = dataset.map(
+            lambda x: {
+                "task_id": x["question_id"],
+                "messages": [{"role": "user", "content": x["question"]}],
+                "test_code": x["test"],
+            },
+            remove_columns=dataset.column_names,
+            num_proc=16,
+        )
+
+    dataset = dataset.map(
+        lambda x: {"task_id": f"{dataset_name}:{x['task_id']}"},
+        remove_columns=["task_id"],
+        num_proc=16,
+    )
+    return dataset
+
+
+def run_distillation(
+    datasets: list,
+    model: str,
+    concurrency: int = 128,
+    eval: str = None,
+    nothink: bool = False,
+    sample_per_prompt: int = 1,
+):
+    if not eval:  # using training data
+        temperature = 0.8
+        output_path = f"{model.split('/')[-1]}.distill.june.train.jsonl"
+        print(f"Expected output path: {output_path}")
+
+        rows = []
+        for args in datasets:
+            print("Loading dataset: ", args)
+            rows.extend([row for row in load_sampled_dataset(*args)])
+        # shuffle for single-turn so that we pack short prompts for faster generation
+        random.seed(666)
+        random.shuffle(rows)
+        dataset = Dataset.from_list(rows, split="train")
+
+        # dataset duplication based on sample_per_prompt
+        if sample_per_prompt > 1:
+            print(
+                f"Duplicating dataset {len(dataset)} times for {sample_per_prompt} samples per prompt"
+            )
+            # localize the duplication to benefit llm cache
+            extended_rows = []
+            for row in dataset:
+                for i in range(sample_per_prompt):
+                    extended_rows.append(
+                        {
+                            "task_id": f"{row['task_id']}@{i}",
+                            "messages": row["messages"],
+                            "test_code": row.get("test_code", None),
+                        }
+                    )
+            dataset = Dataset.from_list(extended_rows, split="train")
+
+    else:  # eval mode
+        assert sample_per_prompt == 1, "Sample per prompt is not supported in eval mode"
+        temperature = 0.2 if "Qwen3" in model or "DeepSeek-R1" in model else 0.0
+        output_path = (
+            f"{model.split('/')[-1]}.distill.june.eval.{eval.split('/')[-1]}.jsonl"
+        )
+        dataset = load_dataset(eval, split="test")
+        # expand turns if there's a mixture of user and assistant turns
+        print(f"Dataset before expansion: {dataset}")
+        expanded_rows = []
+        for row in dataset:
+            if all(m["role"] == "user" for m in row["messages"]):
+                expanded_rows.append(row)
+            else:
+                for i, m in enumerate(row["messages"]):
+                    if m["role"] == "assistant":
+                        expanded_rows.append(
+                            {
+                                "task_id": f"{row['task_id']}:{i-1}",
+                                "messages": row["messages"][:i],
+                            }
+                        )
+        dataset = Dataset.from_list(expanded_rows, split="test")
+        print(f"Dataset after expansion: {dataset}")
+    generated_tasks = set()
+
+    if nothink:
+        output_path = output_path.replace(".jsonl", ".nothink.jsonl")
+
+    print(f"Will process {len(dataset)} samples")
+    print(f"Output path: {output_path}")
+
+    if os.path.exists(output_path):
+        with open(output_path, "r") as f:
+            for line in f:
+                row = json.loads(line)
+                generated_tasks.add(row["task_id"])
+            print(
+                f"Resuming from {output_path}, loaded {len(generated_tasks)} existing results"
+            )
+
+    with ThreadPoolExecutor(max_workers=concurrency) as executor:
+        futures = []
+        for row in dataset:
+            task_id = row["task_id"]
+            if task_id in generated_tasks:
+                continue
+            steps = SAFETY_STEPS
+            if cwe_id := row.get("cwe_id"):
+                steps = SECURITY_STEPS.replace(r"{CWE}", CWE_INFO_COLLECTIONS[cwe_id])
+            elif analyzer_results := row.get("analyzer_results"):
+                steps = SECURITY_STEPS.replace(
+                    r"{CWE}", construction_analyzer_info(analyzer_results)
+                )
+            futures.append(
+                executor.submit(
+                    distillation,
+                    row,
+                    model,
+                    steps=steps,
+                    temperature=temperature,
+                    nothink=nothink,
+                )
+            )
+
+        for future in tqdm(as_completed(futures), total=len(futures)):
+            try:
+                output_row = future.result()
+                with open(output_path, "a") as f:
+                    f.write(json.dumps(output_row, default=str) + "\n")
+                    f.flush()
+            except Exception:
+                print(f"Error:")
+                print_exc()
+                continue
+
+    if eval:
+        eval_main(task=eval, generation_path=output_path)
diff --git a/datagen/ctxdistill/main.py b/datagen/ctxdistill/main.py
new file mode 100644
index 0000000..548db04
--- /dev/null
+++ b/datagen/ctxdistill/main.py
@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from datagen.ctxdistill.ctxdistill import run_distillation
+
+DEFAULT_SAMPLE_RATIO_MEDIUM = 0.25
+
+# Core logic behind:
+# Pick a subset of data {D} total prompts -> {N} prompts
+# Ctx based sampling -> {N * S} where S is the sample size
+# Verfication and pick best of S -> getting {K} prompts with verified responses
+# SFT over {K} prompt-response pairs
+# RL on {D} - {N where most responses are right} prompts, i.e., rm very easy prompts
+
+
+def main(**kwargs):
+    single_turn_datasets = [
+        # mal event / single-turn
+        (
+            "purpcode/mal-event-jailbreak-single-oss-16k",
+            DEFAULT_SAMPLE_RATIO_MEDIUM,
+            4096,
+        ),
+        (
+            "purpcode/mal-event-seed-attack-oss-24k",
+            DEFAULT_SAMPLE_RATIO_MEDIUM,
+            4096,
+        ),
+        # vul code / single-turn
+        ("purpcode/vul2prompt-general-oss-26k", DEFAULT_SAMPLE_RATIO_MEDIUM, 4096),
+        ("purpcode/vul2prompt-benign2vul-oss-21k", DEFAULT_SAMPLE_RATIO_MEDIUM, 4096),
+        ("purpcode/vul2prompt-vul2vul-oss-21k", DEFAULT_SAMPLE_RATIO_MEDIUM, 2048),
+        (
+            "purpcode/vul2prompt-jailbreaking-oss-11k",
+            DEFAULT_SAMPLE_RATIO_MEDIUM,
+            2048,
+        ),
+        # utility
+        ("purpcode/secqa_utility_train", DEFAULT_SAMPLE_RATIO_MEDIUM, 4096),
+        ("KodCode/KodCode-V1-SFT-R1", DEFAULT_SAMPLE_RATIO_MEDIUM, 8192),
+    ]
+    multi_turn_datasets = [
+        # vul code / multi-turn
+        ("purpcode/vul2prompt-multi-oss-5k", 1.0),
+        # mal event / multi-turn
+        ("purpcode/mal-event-fitd-multi-turn-oss-2k", 1.0),
+    ]
+
+    datasets = single_turn_datasets + multi_turn_datasets
+    run_distillation(datasets=datasets, **kwargs)
+
+
+if __name__ == "__main__":
+    from fire import Fire
+
+    Fire(main)
diff --git a/datagen/ctxdistill/post.py b/datagen/ctxdistill/post.py
new file mode 100644
index 0000000..06f608e
--- /dev/null
+++ b/datagen/ctxdistill/post.py
@@ -0,0 +1,357 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Dataset post-processing: Merge conversations, convert file to markdown format, and filter out invalid conversations
+
+import hashlib
+import json
+import os
+from collections import defaultdict
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from copy import deepcopy
+
+import fire
+from datasets import Dataset, DatasetDict
+from markdown_it import MarkdownIt
+from rich import print as rprint
+from tqdm import tqdm
+
+from utils.verify_text import validate_response_structure
+
+
+def _parse_analysis_text(markdown_text: str) -> dict[str, str]:
+    md = MarkdownIt()
+    tokens = md.parse(markdown_text)
+    lines = markdown_text.split("\n")
+
+    # --- Pass 1: Find the minimum header level in the document ---
+    min_level = min(
+        [float("inf")]
+        + [int(token.tag[1:]) for token in tokens if token.type == "heading_open"]
+    )
+    if min_level == float("inf"):
+        return {}
+
+    # --- Pass 2: Extract sections based on the minimum header level found ---
+    top_level_tag = f"h{min_level}"
+    parsed_dict = {}
+    top_level_sections = []
+
+    # Find all header tokens that match the determined top-level tag
+    for i, token in enumerate(tokens):
+        if token.type == "heading_open" and token.tag == top_level_tag:
+            section_title = tokens[i + 1].content
+            # The map gives [start_line, end_line] of the header itself.
+            # Content starts on the line immediately following the header.
+            content_start_line = token.map[1]
+
+            # The section's start line (for finding the end boundary later)
+            section_start_line = token.map[0]
+
+            top_level_sections.append(
+                {
+                    "title": section_title,
+                    "section_start_line": section_start_line,
+                    "content_start_line": content_start_line,
+                }
+            )
+
+    # Determine the content for each top-level section
+    for i, section in enumerate(top_level_sections):
+        # The content for the current section starts after its header
+        content_start = section["content_start_line"]
+
+        # And ends just before the next top-level header begins
+        if (i + 1) == len(top_level_sections):  # is_last_section
+            content_end = len(lines)
+        else:
+            content_end = top_level_sections[i + 1]["section_start_line"]
+
+        content_block = "\n".join(lines[content_start:content_end]).strip()
+        parsed_dict[section["title"]] = content_block
+
+    return {k.lower(): v for k, v in parsed_dict.items()}
+
+
+def transform_md_thinking(response: str) -> str | None:
+    sections = _parse_analysis_text(response)
+    thinking = ["## Analysis"]
+    answer = ["## Answer"]
+    expected_titles = ["intent analysis", "safety analysis", "proposed safe response"]
+    if len(sections) != len(expected_titles):
+        return None
+
+    for i, (title, content) in enumerate(sections.items()):
+        if title != expected_titles[i]:
+            return None
+
+        if i in [0, 1]:
+            thinking.append(content)
+        elif i == 2:
+            if content.startswith('"') and content.endswith('"'):
+                content = content[1:-1]
+            if content.startswith("'") and content.endswith("'"):
+                content = content[1:-1]
+            if content.startswith("```plaintext"):
+                content = content[12:]
+                content = content.removesuffix("```")
+            answer.append(content)
+
+    if "" in thinking:
+        rprint("[red]Error: Empty content in thinking sections[/]")
+        return None
+    if "" in answer:
+        rprint("[red]Error: Empty content in answer sections[/]")
+        return None
+    final_response = "\n".join(thinking) + "\n" + "\n".join(answer)
+    if not validate_response_structure(final_response):
+        rprint("[red]Error: Invalid response structure[/]")
+        return None
+    return final_response
+
+
+def extract_safe_answer(response: str) -> str:
+    response = transform_md_thinking(response) or response
+    return response.split("## Answer")[-1].strip()
+
+
+def _transform_generation(generation):
+    generation = deepcopy(generation)
+    for row in generation["messages"]:
+        if row["role"] == "assistant":
+            if new_content := transform_md_thinking(row["content"]):
+                row["content"] = new_content
+            else:
+                # Signal that this was a "bad conversation"
+                return None
+    return generation
+
+
+def thinking_format(generations):
+    transformed = []
+
+    with ProcessPoolExecutor(max_workers=48) as executor:
+        futures = [
+            executor.submit(_transform_generation, generation)
+            for generation in generations
+        ]
+
+        for future in tqdm(
+            as_completed(futures), desc="Transforming generations", total=len(futures)
+        ):
+            if (generation := future.result()) is not None:
+                transformed.append(generation)
+
+    rprint(
+        f"[green]Transformed {len(transformed)} / {len(generations)} = {len(transformed) / len(generations):.2%} conversations![/]"
+    )
+    return transformed
+
+
+SCORE_THRESHOLD = 0.5
+
+
+def check_rewards(generations, reward_cache_file, batch_size=4096):
+    from rl.grouped_reward import compute_score
+
+    hash2score = {}
+    if os.path.exists(reward_cache_file):
+        with open(reward_cache_file, "r") as f:
+            for line in f:
+                item = json.loads(line)
+                hash2score[item["hash"]] = item["score"]
+        rprint(
+            f"[green]Loaded {len(hash2score)} cached rewards from {reward_cache_file}[/]"
+        )
+
+    hash_fn = lambda gen: hashlib.sha256(
+        json.dumps(gen, sort_keys=True).encode()
+    ).hexdigest()
+    cached_verified_generations = [
+        {**gen, "score": hash2score[hash_fn(gen)]}
+        for gen in generations
+        if hash2score.get(hash_fn(gen), 0) >= SCORE_THRESHOLD
+    ]
+    todo_generations = [gen for gen in generations if hash_fn(gen) not in hash2score]
+
+    batches = [
+        todo_generations[i : i + batch_size]
+        for i in range(0, len(todo_generations), batch_size)
+    ]
+
+    failed_conv_ids = set()
+    verified_generations = []
+    for batch in tqdm(batches):
+        responses = [gen["messages"][-1]["content"] for gen in batch]
+        extra_info = [
+            {"oracles": None, "prompt": gen["messages"][-2]["content"]} for gen in batch
+        ]
+        ground_truth = [None for _ in batch]
+        for i, gen in enumerate(batch):
+            if gen["test_code"]:
+                ground_truth[i] = json.dumps({"pytest": gen["test_code"]})
+                extra_info[i]["oracles"] = ["correctness"]
+            elif "secqa" in gen["task_id"].split(":")[0]:
+                extra_info[i]["oracles"] = ["noreject"]
+            elif "mal-event" in gen["task_id"].split(":")[0]:
+                extra_info[i]["oracles"] = ["general-safety"]
+            elif (
+                "cwe2inst" in gen["task_id"].split(":")[0]
+                or "code2inst" in gen["task_id"].split(":")[0]
+                or "guru2inst" in gen["task_id"].split(":")[0]
+                or "vul-code" in gen["task_id"].split(":")[0]
+                or "vul2prompt" in gen["task_id"].split(":")[0]
+            ):
+                extra_info[i]["oracles"] = ["codesec"]
+            else:
+                raise ValueError(
+                    f"Unknown task_id: {gen['task_id']}, please check the task_id format."
+                )
+
+        scores = compute_score(responses, ground_truth, None, extra_info)
+        for i, score in enumerate(scores):
+            if score >= SCORE_THRESHOLD:
+                verified_generations.append({**batch[i], "score": score})
+            else:
+                failed_conv_ids.add(batch[i]["task_id"].rsplit(":", maxsplit=1)[0])
+            with open(reward_cache_file, "a") as f:
+                f.write(json.dumps({"hash": hash_fn(batch[i]), "score": score}) + "\n")
+
+    # Filter out failed conversations
+    verified_generations = verified_generations + cached_verified_generations
+    n_verified = len(verified_generations)
+    rprint(
+        f"[green]Verified {n_verified} / {len(generations)} = {n_verified / len(generations):.2%} conversations![/]"
+    )
+    return verified_generations
+
+
+def split_into_prefix_subturns(generations):
+    def _split(messages):
+        # assert all user-assistant turns
+        assert len(messages) % 2 == 0
+        assert all(m["role"] in ["assistant", "user"] for m in messages)
+        result = []
+        if len(messages) < 2:
+            return []
+        for i in range(2, len(messages) + 1, 2):
+            if i % 2 == 0:
+                result.append(messages[:i])
+        return result
+
+    new_generations = []
+    for gen in generations:
+        subturns = _split(gen["messages"])
+        for sub in subturns:
+            new_gen = deepcopy(gen)
+            new_gen["messages"] = sub
+            new_gen["task_id"] += f":{len(sub)}"
+            new_generations.append(new_gen)
+    return new_generations
+
+
+def main(generation_path: str, push_to_hub: bool = False):
+    with open(generation_path, "r") as f:
+        generations = [json.loads(line) for line in f]
+
+    sample_per_prompt = 1
+    for gen in generations:
+        if "@" in gen["task_id"]:
+            sample_per_prompt = max(
+                sample_per_prompt, int(gen["task_id"].split("@")[-1])
+            )
+
+    generations = thinking_format(generations)
+    generations = split_into_prefix_subturns(
+        generations
+    )  # necessary as we will cut thinking parts in earilier turns
+    reward_cache_file = generation_path.replace(".jsonl", ".reward_cache.jsonl")
+    print(f"{len(generations)} samples after splitting into prefix subturns")
+    verified_generations = []
+
+    for generation in check_rewards(generations, reward_cache_file=reward_cache_file):
+        generation = deepcopy(generation)
+        # only preserve content and role
+        generation["messages"] = [
+            {"role": m["role"], "content": m["content"], "train": False}
+            for m in generation["messages"]
+        ]
+
+        # pop system prompt
+        if generation["messages"][0]["role"] == "system":
+            generation["messages"].pop(0)
+
+        # set trainable turns and remove thinking parts of past turns
+        for m in generation["messages"][:-1]:
+            if m["role"] == "assistant":
+                m["content"] = m["content"].split("# Answer")[-1].strip()
+        generation["messages"][-1]["train"] = True
+        verified_generations.append(
+            {
+                "task_id": generation["task_id"],
+                "messages": generation["messages"],
+                "score": generation["score"],
+            }
+        )
+
+    # original_task_id -> {sample_id}
+    task_id_to_vsample_ids = defaultdict(list)
+    for vgen in verified_generations:
+        vtask_id = vgen["task_id"]  # {tasl_id}@{sample_id}:{turn_id}
+        task_id = vtask_id.rsplit("@", maxsplit=1)[0]
+        sample_id = vtask_id.rsplit(":", maxsplit=1)[0].rsplit("@", maxsplit=1)[-1]
+        turn_id = vtask_id.rsplit(":", maxsplit=1)[-1]
+        task_id = f"{task_id}:{turn_id}"  # augment task_id with turn_id
+        task_id_to_vsample_ids[task_id].append(
+            {
+                "sample_id": sample_id,
+                "messages": vgen["messages"],
+                "score": vgen["score"],
+            }
+        )
+
+    # write task_ids that have 70%+ verfified samples
+    with open(
+        f"{generation_path.split('/')[-1].split('.distill')[0]}.ez_task_ids.txt", "w"
+    ) as f:
+        for task_id, vsample_ids in task_id_to_vsample_ids.items():
+            if len(vsample_ids) >= sample_per_prompt * 0.7:
+                f.write(task_id + "\n")
+
+    # select one sample per task_id
+    verified_generations = []
+    for task_id, vsamples in task_id_to_vsample_ids.items():
+        selected = vsamples[0]
+        for vsample in vsamples[1:]:
+            if vsample["score"] > selected["score"] or (
+                vsample["score"] == selected["score"]
+                and len(vsample["messages"][-1]["content"])
+                < len(selected["messages"][-1]["content"])
+            ):
+                selected = vsample
+        verified_generations.append(
+            {
+                "task_id": task_id,
+                "messages": selected["messages"],
+                "score": selected["score"],
+            }
+        )
+
+    print(
+        f"Selected {len(verified_generations)} BoN samples from {len(task_id_to_vsample_ids)} task_id:turn"
+    )
+
+    dataset_name = f"ctxdistill-verified-{generation_path.split('/')[-1].split('.distill')[0]}-{len(verified_generations) // 1000}k"
+    with open(f"{dataset_name}.jsonl", "w") as f:
+        for generation in verified_generations:
+            f.write(json.dumps(generation, default=str) + "\n")
+
+    if push_to_hub:
+        dataset = DatasetDict({"train": Dataset.from_list(verified_generations)})
+        dataset.push_to_hub(f"purpcode/{dataset_name}", private=True)
+        rprint(f"[green]Pushed dataset to Hugging Face Hub: {dataset_name}[/]")
+
+
+if __name__ == "__main__":
+    fire.Fire(main)
diff --git a/datagen/rule2code/cwe2code.py b/datagen/rule2code/cwe2code.py
new file mode 100644
index 0000000..287a10f
--- /dev/null
+++ b/datagen/rule2code/cwe2code.py
@@ -0,0 +1,319 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import html
+import json
+import os
+import re
+import textwrap
+import xml.etree.ElementTree as ET
+import zipfile
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from traceback import print_exc
+
+import requests
+from openai import OpenAI
+from rich import print as rprint
+from tqdm import tqdm
+
+
+def transform_xml_to_code(root):
+
+    def get_indent(style):
+        """Return extra indent based on a style like 'margin-left:1em'."""
+        match = re.search(r"margin-left\s*:\s*([\d.]+)em", style)
+        if match:
+            return "  " * int(float(match.group(1)))
+        return ""
+
+    def rec(node, current_indent=""):
+        parts = []
+        if node.text and node.text.strip():
+            text = html.unescape(node.text).strip()
+            parts.append(current_indent + text)
+        for child in node:
+            tag = child.tag.split("}")[-1]
+            if tag == "br":
+                parts.append("\n" + current_indent)
+                if child.tail and child.tail.strip():
+                    parts.append(html.unescape(child.tail).strip())
+            elif tag == "div":
+                extra = get_indent(child.attrib.get("style", ""))
+                new_indent = current_indent + extra
+                child_text = rec(child, new_indent)
+                parts.append("\n" + child_text)
+                if child.tail and child.tail.strip():
+                    parts.append(
+                        "\n" + current_indent + html.unescape(child.tail).strip()
+                    )
+            else:
+                parts.append(rec(child, current_indent))
+        return "".join(parts)
+
+    code = rec(root, "")
+    code = textwrap.dedent(code)
+    lines = code.splitlines()
+    cleaned = "\n".join(line.rstrip() for line in lines if line.strip() != "")
+    return cleaned
+
+
+def parse_examples(examples, ns):
+    if examples is None:
+        return []
+    markdown_output = []
+    examples = [example for example in examples if example is not None]
+    for example in examples:
+        text = f"#### Example {len(markdown_output) + 1}\n"
+        for elem in example:
+            if (
+                elem.tag.endswith("Intro_Text") or elem.tag.endswith("Body_Text")
+            ) and elem.text:
+                text += elem.text.strip() + "\n"
+
+            elif elem.tag.endswith("Example_Code"):
+                language = elem.attrib.get("Language", "").lower()
+                xhtml_div = elem.find("xhtml:div", ns)
+                if xhtml_div is None:
+                    continue
+                block = transform_xml_to_code(xhtml_div)
+                text += f"```{language}\n{block}\n```\n"
+        markdown_output.append(text.strip())
+
+    return markdown_output
+
+
+def create_cwe_information(skip_high_level=True):
+    collection = {}  # return value
+
+    cache_file = "cwec_latest.xml.zip"
+    if os.path.exists(cache_file):
+        print(f"Using cached CWE ZIP file: {cache_file}")
+        z = zipfile.ZipFile(cache_file)
+    else:
+        url = "https://cwe.mitre.org/data/xml/cwec_latest.xml.zip"
+        print(f"Fetching CWE data from: {url}")
+        response = requests.get(url)
+        assert response.status_code == 200, f"Failed to fetch {url}"
+
+        with open(cache_file, "wb") as f:
+            f.write(response.content)
+        print(f"Cached CWE ZIP file to: {cache_file}")
+
+        z = zipfile.ZipFile(cache_file)
+
+    xml_files = [f for f in z.namelist() if f.endswith(".xml")]
+    assert len(xml_files) == 1, "Expected exactly one XML file in the zip"
+    tree = ET.parse(z.open(xml_files[0]))
+
+    root = tree.getroot()
+
+    # Extract the namespace dynamically
+    ns = {
+        "cwe": root.tag.split("}")[0].strip("{"),
+        "xhtml": "http://www.w3.org/1999/xhtml",
+    }
+    weaknesses = root.findall(".//cwe:Weakness", ns)
+    for weakness in weaknesses:
+        markdown_output = []
+        weakness_id = weakness.attrib.get("ID")
+        name = weakness.attrib.get("Name")
+        if weakness.attrib.get("Status") == "Deprecated":
+            continue
+
+        # skip Mapping_Notes - Usage == "Prohibited"
+        if skip_high_level:
+            if (mapping_notes := weakness.find(".//cwe:Mapping_Notes", ns)) is not None:
+                usage = mapping_notes.find(".//cwe:Usage", ns)
+                if usage is not None and usage.text.lower() != "allowed":
+                    print(
+                        f"Skipping CWE-{weakness_id} due to prohibited usage: {usage.text}"
+                    )
+                    continue
+
+        # CWE title
+        markdown_output.append(f"## CWE-{weakness_id}: {name}")
+
+        # Description
+        markdown_output.append("### Description")
+        markdown_output.append(weakness.find(".//cwe:Description", ns).text)
+        edesc = weakness.find(".//cwe:Extended_Description", ns)
+        if edesc is not None and edesc.text and edesc.text.strip():
+            markdown_output.append(edesc.text)
+
+        # Examples
+        examples = weakness.findall(".//cwe:Demonstrative_Example", ns)
+        if expmkd := parse_examples(examples, ns):
+            markdown_output.append("### Examples")
+            markdown_output.extend(expmkd)
+
+        # Notes
+        if (notes := weakness.find(".//cwe:Notes", ns)) is not None:
+            note_text = "### Notes\n"
+            for note in notes:
+                if note.text and note.text.strip():
+                    note_text += "\n* " + note.text.strip()
+            markdown_output.append(note_text)
+
+        # Common_Consequences
+        if (conseq := weakness.find(".//cwe:Common_Consequences", ns)) is not None:
+            conseq_text = ""
+            for csq in conseq:
+                csq = csq.find(".//cwe:Note", ns)
+                if csq is not None and csq.text and csq.text.strip():
+                    conseq_text += "\n* " + csq.text.strip()
+            if conseq_text.strip():
+                markdown_output.append("### Common Consequences\n" + conseq_text)
+
+        markdown_result = "\n\n".join(markdown_output)
+
+        collection[weakness_id] = markdown_result
+
+    return collection
+
+
+def generate_initial_prompt(cwe_id, markdown):
+    return (
+        f"""Given the following markdown text for CWE-{cwe_id}, please demonstrate scenarios (either abstract or realistic) and complete PYTHON code examples that demonstrate this vulnerability:
+
+--- Markdown Content for CWE-{cwe_id} ---
+{markdown}
+--- End of Markdown Content ---
+"""
+        + r"""Now, please create a scenario and PYTHON code example that demonstrates this vulnerability.
+
+Notes:
+* Create a scenario under applicable platforms
+* Format: An example output is demonstrate below:
+
+--- BEGIN OF EXAMPLE ---
+## Code Example
+
+```python
+{code}
+```
+## Explanation
+
+Explain why this code is vulnerable.
+--- END OF EXAMPLE ---"""
+    )
+
+
+def generate_followup_prompt():
+    return r"""Now, please create a NEW scenario and PYTHON code example that demonstrates this vulnerability.
+
+Notes:
+* Creativity & Diversity: Create new scenarios and use different platforms (if applicable)
+* Format: An example output is demonstrate below:
+
+--- BEGIN OF EXAMPLE ---
+## Code Example
+
+```python
+{code}
+```
+## Explanation
+
+Explain why this code is vulnerable.
+--- END OF EXAMPLE ---
+"""
+
+
+def _create_client(remote_api=False):
+    if remote_api:
+        return OpenAI(base_url="https://api.deepseek.com"), "deepseek-reasoner"
+    # Otherwise sglang
+    return OpenAI(api_key="none", base_url="http://0.0.0.0:30000/v1"), "default"
+
+
+def datagen_for_one_cwe(cwe_id, markdown, depth, remote_api=False):
+    assert depth > 0
+
+    client, model = _create_client(remote_api=remote_api)
+    common_args = {"model": model, "temperature": 0.6}
+
+    rprint(f"[bold yellow]Processing: CWE ID: {cwe_id}[/bold yellow]")
+
+    messages = [
+        {
+            "role": "user",
+            "content": generate_initial_prompt(cwe_id, markdown)
+            + generate_followup_prompt(),
+        }
+    ]
+    choice = client.chat.completions.create(messages=messages, **common_args).choices[0]
+
+    if choice.finish_reason == "length":
+        return {"id": cwe_id, "conversation": messages}
+
+    messages.append(
+        {
+            "role": "assistant",
+            "content": choice.message.content.split("</think>")[-1].strip(),
+        }
+    )
+
+    for _ in range(depth - 1):
+        messages.append({"role": "user", "content": generate_followup_prompt()})
+        choice = client.chat.completions.create(
+            messages=messages, **common_args
+        ).choices[0]
+
+        if choice.finish_reason == "length":
+            break
+
+        messages.append(
+            {
+                "role": "assistant",
+                "content": choice.message.content.split("</think>")[-1].strip(),
+            }
+        )
+
+    rprint(f"[bold green]Completed: CWE ID: {cwe_id}[/bold green]")
+    return {"id": cwe_id, "conversation": messages}
+
+
+def main(
+    parallel=256,
+    output_path="outputs/rule2code/cwe2code-raw.jsonl",
+    depth=1,
+    remote_api=False,
+):
+    collection = create_cwe_information()
+    # each line: cwe_id, conversation
+
+    finished = set()
+    if os.path.exists(output_path):
+        with open(output_path, "r", encoding="utf-8") as f:
+            for line in f:
+                finished.add(json.loads(line)["id"])
+
+    with ThreadPoolExecutor(max_workers=parallel) as executor:
+        futures = []
+        for cwe_id, markdown in collection.items():
+            if cwe_id in finished:
+                continue
+            futures.append(
+                executor.submit(
+                    datagen_for_one_cwe, cwe_id, markdown, depth, remote_api
+                )
+            )
+
+        for future in tqdm(as_completed(futures), total=len(futures)):
+            result = future.result()
+            try:
+                result_json = json.dumps(result, ensure_ascii=False)
+            except Exception:
+                print("Error found -- ")
+                print(f"{result = }")
+                print_exc()
+                continue
+
+            with open(output_path, "a", encoding="utf-8") as f:
+                f.write(result_json + "\n")
+
+
+if __name__ == "__main__":
+    from fire import Fire
+
+    Fire(main)
diff --git a/datagen/rule2code/get_bandit_rules.py b/datagen/rule2code/get_bandit_rules.py
new file mode 100644
index 0000000..a1ec64d
--- /dev/null
+++ b/datagen/rule2code/get_bandit_rules.py
@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# TODO(@zhewang2001): Please refactor the corresponding code snippets and then upload it.
diff --git a/datagen/rule2code/guru2code.py b/datagen/rule2code/guru2code.py
new file mode 100644
index 0000000..a1ec64d
--- /dev/null
+++ b/datagen/rule2code/guru2code.py
@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# TODO(@zhewang2001): Please refactor the corresponding code snippets and then upload it.
diff --git a/datagen/rule2code/post_process.py b/datagen/rule2code/post_process.py
new file mode 100644
index 0000000..a1ec64d
--- /dev/null
+++ b/datagen/rule2code/post_process.py
@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# TODO(@zhewang2001): Please refactor the corresponding code snippets and then upload it.
diff --git a/sft/controlled/data_split_for_data_ablation.py b/sft/controlled/data_split_for_data_ablation.py
index e698324..6b1dffb 100644
--- a/sft/controlled/data_split_for_data_ablation.py
+++ b/sft/controlled/data_split_for_data_ablation.py
@@ -19,14 +19,14 @@
     for r in data
     if r["task_id"].split(":")[0]
     in [
-        "purpcorn/vul2prompt-jailbreaking-oss-11k",
-        "purpcorn/vul2prompt-benign2vul-oss",
-        "purpcorn/mal-event-fitd-multi-turn-oss-2k",
-        "purpcorn/vul2prompt-multi-oss-5k",
-        "purpcorn/mal-event-seed-attack-oss-24k",
-        "purpcorn/mal-event-jailbreak-single-oss-16k",
-        "purpcorn/vul2prompt-general-oss",
-        "purpcorn/vul2prompt-vul2vul-oss",
+        "purpcode/vul2prompt-jailbreaking-oss-11k",
+        "purpcode/vul2prompt-benign2vul-oss-21k",
+        "purpcode/mal-event-fitd-multi-turn-oss-2k",
+        "purpcode/vul2prompt-multi-oss-5k",
+        "purpcode/mal-event-seed-attack-oss-24k",
+        "purpcode/mal-event-jailbreak-single-oss-16k",
+        "purpcode/vul2prompt-general-oss-26k",
+        "purpcode/vul2prompt-vul2vul-oss-21k",
     ]
 ]
 
@@ -35,7 +35,7 @@
     for r in data
     if r["task_id"].split(":")[0]
     in [
-        "purpcorn/secqa_utility_train",
+        "purpcode/secqa_utility_train",
         "KodCode/KodCode-V1-SFT-R1",
     ]
 ]
diff --git a/utils/verify_text.py b/utils/verify_text.py
new file mode 100644
index 0000000..bf4c8b3
--- /dev/null
+++ b/utils/verify_text.py
@@ -0,0 +1,108 @@
+# SPDX-FileCopyrightText: (c) UIUC PurpCode Team
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from typing import Tuple
+
+# Legacy thinking format based on XML
+XML_THINK_PATTERN = re.compile(r"<think>(.{16,})</think>.*", re.DOTALL)
+
+
+def xml_validate_response_structure(response: str) -> bool:
+    if match := XML_THINK_PATTERN.match(response):
+        thinking = f"<think>{match.group(1)}</think>"
+        response = (
+            response.split(thinking)[-1]
+            .strip()
+            .removeprefix("<answer>")
+            .removesuffix("</answer>")
+            .strip()
+        )
+        # <answer>, </answer>, <think>, </think> should not be in the response
+        return (
+            "<answer>" not in response
+            and "</answer>" not in response
+            and "<think>" not in response
+            and "</think>" not in response
+        )
+
+    return False
+
+
+def try_extract_answer_from_xml(response: str) -> str:
+    answer_pattern = r"<think>(.*?)</think>"
+    if matches := list(re.finditer(answer_pattern, response, re.DOTALL)):
+        thinking = f"<think>{matches[-1].group(1)}</think>"
+        response = response.split(thinking)[-1].strip()
+    return response.removeprefix("<answer>").removesuffix("</answer>").strip()
+
+
+# Thinking format based on Markdown
+MD_THINK_PATTERN = re.compile(r"^## Analysis\n(.{16,})\n## Answer\n(.{8,})$", re.DOTALL)
+
+
+def markdown_validate_response_structure(response: str) -> bool:
+    prefill = "## Analysis\n"
+    response = prefill + response.strip().lstrip(prefill)
+    return bool(MD_THINK_PATTERN.match(response.strip()))
+
+
+def try_extract_answer_from_markdown(response: str) -> str:
+    return response.split("\n## Answer", maxsplit=1)[-1].strip()
+
+
+# Select mode
+validate_response_structure = markdown_validate_response_structure
+try_extract_answer = try_extract_answer_from_markdown
+
+
+def _has_repetition(s: str, rep_length_thresh: int = 32, rep_count_thresh: int = 8):
+    if not s or len(s) < rep_length_thresh:
+        return False, ""
+
+    subsequence_count = {}
+    length = rep_length_thresh
+
+    for i in range(len(s) - length + 1):
+        subseq = s[i : i + length]
+        subsequence_count[subseq] = subsequence_count.get(subseq, 0) + 1
+        if subsequence_count[subseq] >= rep_count_thresh:
+            return (
+                True,
+                "-" * 16
+                + "Repitition Check"
+                + "-" * 16
+                + f"\nRepeated {subsequence_count[subseq]} times: {subseq}",
+            )
+
+    return False, ""
+
+
+CODE_PATTERN = re.compile(r"```(?:\w+)?\n(.*?)\n```", re.DOTALL)
+
+
+def extract_code_from_string(solution_str):
+    code_blocks = CODE_PATTERN.findall(solution_str)
+    return "\n".join(code_blocks).strip()
+
+
+def check_fmt(response: str) -> Tuple[bool, str]:
+    reward_log = (
+        "-" * 16
+        + "Bad format detected -- Original Model Output"
+        + "-" * 16
+        + "\n"
+        + response
+    )
+
+    if not validate_response_structure(response):
+        return False, reward_log
+
+    if (rep_res := _has_repetition(response))[0]:
+        return False, rep_res[1]
+
+    if len(try_extract_answer(response)) == 0:
+        return False, reward_log
+
+    return True, ""