Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+wikilarge-dataset/wiki.full.aner.ori.train.95.tsv filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,9 @@
+.venv/*
+bigcode_humanevalpack_full/*
+google_code_x_glue_ct_code_to_text_full/*
+wiki_results.tsv
+.cache/*
+checkpoints/*
+code_x_glue_cc_code_refinement_full/*
+wikilarge-dataset/*
+

=0.14.0

@@ -0,0 +1,18 @@
+Looking in indexes: https://bytedpypi.byted.org/simple/
+Requirement already satisfied: huggingface-hub in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (1.2.3)
+Requirement already satisfied: filelock in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (3.20.0)
+Requirement already satisfied: fsspec>=2023.5.0 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (2025.10.0)
+Requirement already satisfied: hf-xet<2.0.0,>=1.2.0 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (1.2.0)
+Requirement already satisfied: httpx<1,>=0.23.0 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (0.28.1)
+Requirement already satisfied: packaging>=20.9 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (25.0)
+Requirement already satisfied: pyyaml>=5.1 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (6.0.3)
+Requirement already satisfied: shellingham in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (1.5.4)
+Requirement already satisfied: tqdm>=4.42.1 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (4.67.1)
+Requirement already satisfied: typer-slim in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (0.20.0)
+Requirement already satisfied: typing-extensions>=3.7.4.3 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from huggingface-hub) (4.15.0)
+Requirement already satisfied: anyio in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface-hub) (4.12.0)
+Requirement already satisfied: certifi in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface-hub) (2025.11.12)
+Requirement already satisfied: httpcore==1.* in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface-hub) (1.0.9)
+Requirement already satisfied: idna in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface-hub) (3.11)
+Requirement already satisfied: h11>=0.16 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->huggingface-hub) (0.16.0)
+Requirement already satisfied: click>=8.0.0 in /mlx_devbox/users/lixinyu.222/playground/Diffusion_Learning/.venv/lib/python3.11/site-packages (from typer-slim->huggingface-hub) (8.3.1)

README.md

@@ -0,0 +1,49 @@
+# CodeFlow
+## Version 1
+CodeFlow/
+├── src/
+│   ├── __init__.py
+│   ├── config.py          # 全局配置
+│   ├── models/
+│   │   ├── __init__.py
+│   │   ├── autoencoder.py # 潜空间 AE
+│   │   └── dit.py         # Diffusion Transformer
+│   ├── utils/
+│   │   ├── sandbox.py     # 代码执行沙箱
+│   │   └── data_utils.py  # 数据加载器
+│   └── trainer.py         # 训练与推理引擎
+├── tests/
+│   ├── test_models.py     # 模型单元测试
+│   └── test_sandbox.py    # 沙箱单元测试
+├── run_wiki_flow.py       # 入口1：Wiki 简化
+└── run_mbpp_ae.py         # 入口2：MBPP 重建验证
+
+## Version 2
+CodeFlow/
+├── src/
+│   ├── __init__.py
+│   ├── config.py          # 全局配置 (Patching, Dimensions)
+│   ├── models/
+│   │   ├── __init__.py
+│   │   ├── autoencoder.py # Jina -> Linear -> Sphere -> Decoder
+│   │   └── dit.py         # Patched DiT + Flow Logic
+│   ├── utils/
+│   │   └── data.py        # Wiki/MBPP 数据加载
+│   └── trainer.py         # 训练引擎 (AE & Flow)
+├── run_mbpp_ae.py         # 入口1：验证重建能力
+├── run_wiki_flow.py       # 入口2：验证 Flow Matching 编辑能力
+└── requirements.txt
+
+Autoencoder: 移除 VAE/KL，改用 Linear Compression + L2 Normalization。这保证了潜空间在单位球面上，语义连续且训练极其稳定。
+Backbone: 依然是 Jina-v2 (Freeze) + NAR Decoder。
+Generator: Patched DiT 配合 Rectified Flow，解决长序列计算瓶颈。
+Optimization: 内置梯度累积、混合精度开关（默认关闭以适配 Jina）、多进程数据处理。
+
+### 手动下载
+# 安装 huggingface-hub 工具（若未安装）
+pip install huggingface-hub
+
+# 下载模型到本地目录（比如 ./jina-embeddings-v2-base-code）
+huggingface-cli download --resume-download jinaai/jina-embeddings-v2-base-code --local-dir ./jina-embeddings-v2-base-code
+
+huggingface-cli download bogdancazan/wikilarge-text-simplification --repo-type dataset --resume-download --local-dir ./wikilarge-dataset

checkpoints/ae_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f586f7c121051addeac0aa2e324cef974a2c46b4737c82d7b1835411efdca56d
+size 1481302557

checkpoints/ae_last.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55ed2a5d4a6e5c47fa7e148da7d296096720b368947b54b72ca2d487d6581116
+size 1481302557

checkpoints/flow_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:909a902719e073cdb523c84cab003e8df3cb0763db4b78e38b6e5a08c529f193
+size 532346699

checkpoints/flow_last.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39d42298d1e93ead94725f2585425e8bdbc5d196dccaba57348b9bdc12ee56a1
+size 532346699

code_x_glue_cc_code_refinement_full/.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

code_x_glue_cc_code_refinement_full/README.md

@@ -0,0 +1,273 @@
+---
+annotations_creators:
+- expert-generated
+language_creators:
+- found
+language:
+- code
+license:
+- c-uda
+multilinguality:
+- other-programming-languages
+size_categories:
+- 10K<n<100K
+source_datasets:
+- original
+task_categories:
+- text2text-generation
+task_ids: []
+pretty_name: CodeXGlueCcCodeRefinement
+tags:
+- debugging
+dataset_info:
+- config_name: medium
+  features:
+  - name: id
+    dtype: int32
+  - name: buggy
+    dtype: string
+  - name: fixed
+    dtype: string
+  splits:
+  - name: train
+    num_bytes: 32614786
+    num_examples: 52364
+  - name: validation
+    num_bytes: 4086733
+    num_examples: 6546
+  - name: test
+    num_bytes: 4063665
+    num_examples: 6545
+  download_size: 14929559
+  dataset_size: 40765184
+- config_name: small
+  features:
+  - name: id
+    dtype: int32
+  - name: buggy
+    dtype: string
+  - name: fixed
+    dtype: string
+  splits:
+  - name: train
+    num_bytes: 13006679
+    num_examples: 46680
+  - name: validation
+    num_bytes: 1629242
+    num_examples: 5835
+  - name: test
+    num_bytes: 1619700
+    num_examples: 5835
+  download_size: 5894462
+  dataset_size: 16255621
+configs:
+- config_name: medium
+  data_files:
+  - split: train
+    path: medium/train-*
+  - split: validation
+    path: medium/validation-*
+  - split: test
+    path: medium/test-*
+- config_name: small
+  data_files:
+  - split: train
+    path: small/train-*
+  - split: validation
+    path: small/validation-*
+  - split: test
+    path: small/test-*
+---
+
+# Dataset Card for "code_x_glue_cc_code_refinement"
+
+## Table of Contents
+- [Dataset Description](#dataset-description)
+  - [Dataset Summary](#dataset-summary)
+  - [Supported Tasks and Leaderboards](#supported-tasks)
+  - [Languages](#languages)
+- [Dataset Structure](#dataset-structure)
+  - [Data Instances](#data-instances)
+  - [Data Fields](#data-fields)
+  - [Data Splits](#data-splits-sample-size)
+- [Dataset Creation](#dataset-creation)
+  - [Curation Rationale](#curation-rationale)
+  - [Source Data](#source-data)
+  - [Annotations](#annotations)
+  - [Personal and Sensitive Information](#personal-and-sensitive-information)
+- [Considerations for Using the Data](#considerations-for-using-the-data)
+  - [Social Impact of Dataset](#social-impact-of-dataset)
+  - [Discussion of Biases](#discussion-of-biases)
+  - [Other Known Limitations](#other-known-limitations)
+- [Additional Information](#additional-information)
+  - [Dataset Curators](#dataset-curators)
+  - [Licensing Information](#licensing-information)
+  - [Citation Information](#citation-information)
+  - [Contributions](#contributions)
+
+## Dataset Description
+
+- **Homepage:** https://github.com/microsoft/CodeXGLUE/tree/main/Code-Code/code-refinement
+- **Paper:** https://arxiv.org/abs/2102.04664
+
+### Dataset Summary
+
+CodeXGLUE code-refinement dataset, available at https://github.com/microsoft/CodeXGLUE/tree/main/Code-Code/code-refinement
+
+We use the dataset released by this paper(https://arxiv.org/pdf/1812.08693.pdf). The source side is a Java function with bugs and the target side is the refined one. All the function and variable names are normalized. Their dataset contains two subsets ( i.e.small and medium) based on the function length.
+
+### Supported Tasks and Leaderboards
+
+- `text2text-generation-other-debugging`: The dataset can be used to train a model for automatically fixing buggy code.
+
+### Languages
+
+- Java **programming** language
+
+## Dataset Structure
+
+### Data Instances
+
+#### medium
+
+An example of 'train' looks as follows.
+```
+{
+    "buggy": "public static TYPE_1 init ( java.lang.String name , java.util.Date date ) { TYPE_1 VAR_1 = new TYPE_1 ( ) ; VAR_1 . METHOD_1 ( name ) ; java.util.Calendar VAR_2 = java.util.Calendar.getInstance ( ) ; VAR_2 . METHOD_2 ( date ) ; VAR_1 . METHOD_3 ( VAR_2 ) ; return VAR_1 ; }\n", 
+    "fixed": "public static TYPE_1 init ( java.lang.String name , java.util.Date date ) { TYPE_1 VAR_1 = new TYPE_1 ( ) ; VAR_1 . METHOD_1 ( name ) ; java.util.Calendar VAR_2 = null ; if ( date != null ) { VAR_2 = java.util.Calendar.getInstance ( ) ; VAR_2 . METHOD_2 ( date ) ; } VAR_1 . METHOD_3 ( VAR_2 ) ; return VAR_1 ; }\n", 
+    "id": 0
+}
+```
+
+#### small
+
+An example of 'validation' looks as follows.
+```
+{
+    "buggy": "public java.util.List < TYPE_1 > METHOD_1 ( ) { java.util.ArrayList < TYPE_1 > VAR_1 = new java.util.ArrayList < TYPE_1 > ( ) ; for ( TYPE_2 VAR_2 : VAR_3 ) { VAR_1 . METHOD_2 ( VAR_2 . METHOD_1 ( ) ) ; } return VAR_1 ; } \n", 
+    "fixed": "public java.util.List < TYPE_1 > METHOD_1 ( ) { return VAR_1 ; } \n", 
+    "id": 0
+}
+```
+
+### Data Fields
+
+In the following each data field in go is explained for each config. The data fields are the same among all splits.
+
+#### medium, small
+
+|field name| type |          description           |
+|----------|------|--------------------------------|
+|id        |int32 | Index of the sample            |
+|buggy     |string| The buggy version of the code  |
+|fixed     |string| The correct version of the code|
+
+### Data Splits
+
+| name |train|validation|test|
+|------|----:|---------:|---:|
+|medium|52364|      6546|6545|
+|small |46680|      5835|5835|
+
+## Dataset Creation
+
+### Curation Rationale
+
+[More Information Needed]
+
+### Source Data
+
+#### Initial Data Collection and Normalization
+
+Downloaded from GitHub Archive every public GitHub event between March 2011 and October 2017 and used the Google BigQuery APIs.
+[More Information Needed]
+
+#### Who are the source language producers?
+
+Software Engineering developers.
+
+### Annotations
+
+#### Annotation process
+
+Automatically annotated by filtering commit messages containing the pattern: ("fix" or "solve") and ("bug" or "issue" or "problem" or "error"). A statistically significant amount of samples (95% confidence level with 5% confidence interval) were manually evaluated by two authors to check if the filtered bug/fix pairs were correct. After all disagreements were settled, authors conclude that 97.6% were true positives.
+
+#### Who are the annotators?
+
+Heuristics and the authors of the paper.
+
+### Personal and Sensitive Information
+
+[More Information Needed]
+
+## Considerations for Using the Data
+
+### Social Impact of Dataset
+
+[More Information Needed]
+
+### Discussion of Biases
+
+[More Information Needed]
+
+### Other Known Limitations
+
+[More Information Needed]
+
+## Additional Information
+
+### Dataset Curators
+
+https://github.com/microsoft, https://github.com/madlag
+
+### Licensing Information
+
+Computational Use of Data Agreement (C-UDA) License.
+
+### Citation Information
+
+```
+@article{DBLP:journals/corr/abs-2102-04664,
+  author    = {Shuai Lu and
+               Daya Guo and
+               Shuo Ren and
+               Junjie Huang and
+               Alexey Svyatkovskiy and
+               Ambrosio Blanco and
+               Colin B. Clement and
+               Dawn Drain and
+               Daxin Jiang and
+               Duyu Tang and
+               Ge Li and
+               Lidong Zhou and
+               Linjun Shou and
+               Long Zhou and
+               Michele Tufano and
+               Ming Gong and
+               Ming Zhou and
+               Nan Duan and
+               Neel Sundaresan and
+               Shao Kun Deng and
+               Shengyu Fu and
+               Shujie Liu},
+  title     = {CodeXGLUE: {A} Machine Learning Benchmark Dataset for Code Understanding
+               and Generation},
+  journal   = {CoRR},
+  volume    = {abs/2102.04664},
+  year      = {2021}
+}
+@article{tufano2019empirical,
+  title={An empirical study on learning bug-fixing patches in the wild via neural machine translation},
+  author={Tufano, Michele and Watson, Cody and Bavota, Gabriele and Penta, Massimiliano Di and White, Martin and Poshyvanyk, Denys},
+  journal={ACM Transactions on Software Engineering and Methodology (TOSEM)},
+  volume={28},
+  number={4},
+  pages={1--29},
+  year={2019},
+  publisher={ACM New York, NY, USA}
+}
+```
+
+### Contributions
+
+Thanks to @madlag (and partly also @ncoop57) for adding this dataset.

code_x_glue_cc_code_refinement_full/medium/test-00000-of-00001.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:776766b34878a999193dae39f0462ed1dc5aec3b7d219d3488db4c0127eca858
+size 1488083

code_x_glue_cc_code_refinement_full/medium/train-00000-of-00001.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:955496550f8603dd920aa2a09d1e4a14d878bda9f35f2f3575bae68b9d9493f0
+size 11943277

code_x_glue_cc_code_refinement_full/medium/validation-00000-of-00001.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9eb660729a7a43df19c714726933a517d0497cefa99cb95328bc964658eaa44c
+size 1498199

code_x_glue_cc_code_refinement_full/small/test-00000-of-00001.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be9b78909122cc07a9b31b55c999fe8b8346b1bd775e726b0cbb044f9d1ebc90
+size 588578

code_x_glue_cc_code_refinement_full/small/train-00000-of-00001.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e0ddaf02afd36f2827d0f6edcf94b3b6b5908206cbf63fc6ac99ee4a6332472
+size 4715251

code_x_glue_cc_code_refinement_full/small/validation-00000-of-00001.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b623a6544959b1767eb3082064b3b3896e26f0770dfd2270d14d4e62da0246a5
+size 590633

download_data.txt

@@ -0,0 +1,29 @@
+
+pip install huggingface-hub
+# 升级到最新版本（推荐）
+pip install --upgrade huggingface-hub
+
+
+# 若需指定稳定版本（可选，最低支持--subset的版本约为0.14.0）
+pip install huggingface-hub>=0.14.0
+
+### 烙铁下错数据集了
+huggingface-cli download \
+  google/code_x_glue_ct_code_to_text \
+  --repo-type dataset \
+  --local-dir ./google_code_x_glue_ct_code_to_text_full \
+  --local-dir-use-symlinks False
+
+huggingface-cli download \
+  bigcode/humanevalpack \
+  --repo-type dataset \
+  --local-dir ./bigcode_humanevalpack_full \
+  --local-dir-use-symlinks False
+
+## code Refine 数据集
+load_dataset("google/code_x_glue_cc_code_refinement", "medium")
+huggingface-cli download \
+  google/code_x_glue_cc_code_refinement \
+  --repo-type dataset \
+  --local-dir ./code_x_glue_cc_code_refinement_full \
+  --local-dir-use-symlinks False

eval_ae.py

@@ -0,0 +1,156 @@
+# scripts/eval_ae_consistency.py
+"""
+z0 = encoder(x)
+x^1 = decoder(z0)
+z1 = encoder(x^1)
+"""
+
+import argparse
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+from src.config import ModelConfig, TrainConfig
+from src.models.autoencoder import ReshapedAutoencoder
+from src.utils.data_utils import prepare_data
+
+def pick_stop_id(tokenizer):
+    return tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.sep_token_id
+
+def masked_mean(x, mask, eps=1e-6):
+    # x: [B,L] or [B,L,D] reduced already, mask: [B,L]
+    denom = mask.sum().clamp(min=eps)
+    return (x * mask).sum() / denom
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", type=str, default="wiki")
+    ap.add_argument("--split", type=str, default="test")
+    ap.add_argument("--max_seq_len", type=int, default=128)
+    ap.add_argument("--batch_size", type=int, default=16)
+    ap.add_argument("--ckpt", type=str, default="/mnt/hdfs/user/lixinyu.222/CodeFlow/residual_robust_checkpoints/ae_best.pt", help="path to ae.state_dict()")
+    ap.add_argument("--max_batches", type=int, default=0, help="0 means full eval")
+    ap.add_argument("--print_n", type=int, default=8)
+    args = ap.parse_args()
+
+    # configs
+    m_cfg = ModelConfig(
+        encoder_name='../jina-embeddings-v2-base-code',
+        latent_dim=512,
+        max_seq_len=args.max_seq_len,
+    )
+    t_cfg = TrainConfig(batch_size=args.batch_size)
+
+    device = t_cfg.device
+    tokenizer = AutoTokenizer.from_pretrained(m_cfg.encoder_name,local_files_only=True,trust_remote_code=False)
+    stop_id = pick_stop_id(tokenizer)
+
+    loader = prepare_data(args.dataset, tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split=args.split)
+    #  test_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="test")
+
+    ae = ReshapedAutoencoder(m_cfg).to(device).float()
+    if args.ckpt:
+        sd = torch.load(args.ckpt, map_location="cpu")
+        ae.load_state_dict(sd, strict=True)
+    ae.eval()
+
+    total_ce = 0.0
+    total_acc = 0.0
+    total_tokens = 0.0
+
+    eos_found = 0
+    eos_pos_err = 0.0
+    eos_count = 0
+
+    total_cos = 0.0
+    total_l2 = 0.0
+    total_lat_tokens = 0.0
+
+    printed = 0
+
+    for bi, batch in enumerate(tqdm(loader, desc="Eval AE")):
+        if args.max_batches and bi >= args.max_batches:
+            break
+
+        ids = batch["tgt_ids"].to(device)
+        mask = batch["tgt_mask"].to(device)
+
+        # --- forward ---
+        z0 = ae.encode(ids, mask)                                  # [B,L,D]
+        logits = ae.decode(z0, attention_mask=mask)                # [B,L,V]
+        pred = logits.argmax(dim=-1)                               # [B,L]
+
+        # --- masked CE ---
+        labels = ids.masked_fill(mask == 0, -100)
+        ce = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100, reduction="sum")
+        total_ce += ce.item()
+
+        # --- token acc (masked) ---
+        correct = ((pred == ids) & (mask.bool())).sum().item()
+        tok = mask.sum().item()
+        total_acc += correct
+        total_tokens += tok
+
+        # --- EOS stats ---
+        # true/pred EOS position (first occurrence)
+        B, L = ids.shape
+        for i in range(B):
+            # only search within valid tokens
+            valid_len = int(mask[i].sum().item())
+            true_seq = ids[i, :valid_len]
+            pred_seq = pred[i, :valid_len]
+
+            true_pos = (true_seq == stop_id).nonzero(as_tuple=True)[0]
+            pred_pos = (pred_seq == stop_id).nonzero(as_tuple=True)[0]
+
+            if pred_pos.numel() > 0:
+                eos_found += 1
+            if true_pos.numel() > 0:
+                eos_count += 1
+                tpos = int(true_pos[0].item())
+                ppos = int(pred_pos[0].item()) if pred_pos.numel() > 0 else valid_len - 1
+                eos_pos_err += abs(ppos - tpos)
+
+        # --- latent cycle: z0 -> token -> z1 ---
+        z1 = ae.encode(pred, mask)
+        cos = F.cosine_similarity(z0, z1, dim=-1)                  # [B,L]
+        l2 = (z0 - z1).pow(2).mean(dim=-1)                         # [B,L]
+        total_cos += (cos * mask).sum().item()
+        total_l2 += (l2 * mask).sum().item()
+        total_lat_tokens += mask.sum().item()
+
+        # --- print a few examples ---
+        if printed < args.print_n:
+            s = tokenizer.decode(ids[0], skip_special_tokens=True)
+            ## 这里没有进行 pos 截断
+            # valid_len = int(mask[0].sum().item())
+            # pred_seq = pred[0, :valid_len]
+            # # 找 stop（eos/sep）
+            # end = _first_pos(pred_seq, stop_id, default=valid_len-1) + 1
+            # g = tokenizer.decode(pred_seq[:end], skip_special_tokens=True)
+            g = tokenizer.decode(pred[0], skip_special_tokens=True)
+            print("\n--- Example ---")
+            print("GT :", s)
+            print("REC:", g)
+            printed += 1
+
+    avg_ce = total_ce / max(total_tokens, 1.0)
+    avg_acc = total_acc / max(total_tokens, 1.0)
+    avg_cos = total_cos / max(total_lat_tokens, 1.0)
+    avg_l2 = total_l2 / max(total_lat_tokens, 1.0)
+
+    eos_found_rate = eos_found / max(total_tokens / args.max_seq_len, 1.0)  # 近似 batch 数
+    eos_mae = eos_pos_err / max(eos_count, 1)
+
+    print("\n===== AE Metrics =====")
+    print(f"Masked CE per-token: {avg_ce:.4f}")
+    print(f"Token Acc (masked): {avg_acc:.4f}")
+    print(f"Latent cycle cosine(z0,z1): {avg_cos:.4f}")
+    print(f"Latent cycle l2(z0,z1): {avg_l2:.6f}")
+    print(f"EOS found rate (rough): {eos_found_rate:.4f}")
+    print(f"EOS position MAE (only where GT has EOS): {eos_mae:.2f}")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,9 @@
+torch 
+transformers 
+datasets 
+accelerate 
+scikit-learn
+timm
+evaluate
+sacrebleu
+sacremoses

run_repair_flow.py

@@ -0,0 +1,162 @@
+import torch
+import torch.optim as optim
+from transformers import AutoTokenizer
+from tqdm import tqdm
+
+from src.config import ModelConfig, TrainConfig
+from src.models.autoencoder import SphericalAutoencoder
+from src.models.dit import PatchedFlowDiT
+from src.trainer import Trainer
+from src.utils.data_utils import prepare_data
+from src.utils.sandbox import SafeSandbox
+from src.search import DiffuMCTS
+
+def inference(ae, flow, src_ids, src_mask, device, steps=10):
+    ae.eval(); flow.eval()
+    with torch.no_grad():
+        # Encode Source (Buggy) -> z_0
+        z_curr = ae.encode(src_ids, src_mask)
+        z_cond = z_curr.clone()
+        
+        dt = 1.0 / steps
+        for i in range(steps):
+            t = torch.ones(z_curr.shape[0], device=device) * (i / steps)
+            v = flow(z_curr, t, condition=z_cond).float()
+            z_curr = z_curr + v * dt
+            
+        z_curr = torch.nn.functional.normalize(z_curr, p=2, dim=-1)
+        logits = ae.decode(z_curr)
+        return torch.argmax(logits, dim=-1)
+
+def evaluate_on_humaneval(ae, flow, tokenizer, device, num_samples=20):
+    """
+    在 HumanEvalPack 上进行真实的执行测试
+    """
+    print("\n>>> Starting Evaluation on HumanEvalPack (Real Execution)...")
+    loader = prepare_data("humanevalpack", tokenizer, 512, 1, split="test")
+    sandbox = SafeSandbox()
+    
+    passed = 0
+    total = 0
+    
+    # 只测前 num_samples 个，节省时间
+    for i, batch in enumerate(tqdm(loader, total=num_samples)):
+        if i >= num_samples: break
+        
+        src = batch['src_ids'].to(device)
+        mask = batch['src_mask'].to(device)
+        test_code = batch['test_code'][0]
+        entry_point = batch['entry_point'][0]
+        
+        # 1. Flow Inference
+        out_ids = inference(ae, flow, src, mask, device)
+        gen_code = tokenizer.decode(out_ids[0], skip_special_tokens=True)
+        
+        # 2. Sandbox Execution
+        is_pass, msg = sandbox.run(gen_code, test_code, entry_point)
+        
+        if is_pass:
+            passed += 1
+        
+        total += 1
+        
+        # 打印第一个 Case 看看效果
+        if i == 0:
+            print(f"\n[Case 0] Pass: {is_pass}")
+            print(f"Error: {msg}")
+            print(f"Generated:\n{gen_code[:200]}...")
+
+    print(f"\n=== Eval Result ===")
+    print(f"Pass@1: {passed}/{total} = {passed/total*100:.2f}%")
+
+def evaluate_with_mcts(ae, flow, tokenizer, device, num_samples=20):
+    """
+    使用 Diffu-MCTS 进行强化评估
+    """
+    print(f"\n>>> Starting Diffu-MCTS Evaluation (samples={num_samples})...")
+    
+    # 1. 准备数据和组件
+    loader = prepare_data("humanevalpack", tokenizer, 512, 1, split="test")
+    sandbox = SafeSandbox(timeout=2.0) # 2秒超时防止死循环
+    
+    # 2. 初始化搜索器
+    mcts = DiffuMCTS(ae, flow, tokenizer, sandbox, device, config=None)
+    mcts.num_branches = 8 # 设定分支数 K=8
+    
+    passed = 0
+    total = 0
+    
+    # 3. 评估循环
+    for i, batch in enumerate(tqdm(loader, total=num_samples)):
+        if i >= num_samples: break
+        
+        # 提取原始文本 (因为 MCTS 内部会处理 Tokenize)
+        # 这里的 batch['src_ids'] 是 tensor，我们需要原始 string
+        # 但 data_loader 把 string 丢了，所以我们这里反解码一下，或者修改 prepare_data 返回 raw text
+        # 为了简单，我们反解码 Buggy Code
+        src_ids = batch['src_ids'].to(device)
+        buggy_code = tokenizer.decode(src_ids[0], skip_special_tokens=True)
+        
+        test_code = batch['test_code'][0]
+        entry_point = batch['entry_point'][0]
+        
+        # --- 调用 MCTS ---
+        fixed_code, is_success = mcts.solve(buggy_code, test_code, entry_point)
+        
+        if is_success:
+            passed += 1
+        
+        total += 1
+        
+        # Log 第一个样本
+        if i == 0:
+            print(f"\n[Case 0]")
+            print(f"Buggy:\n{buggy_code[:100]}...")
+            print(f"Fixed:\n{fixed_code[:100]}...")
+            print(f"Result: {'✅ PASS' if is_success else '❌ FAIL'}")
+
+    print(f"\n=== MCTS Results ===")
+    print(f"Pass@1 (with Search K={mcts.num_branches}): {passed}/{total} = {passed/total*100:.2f}%")
+
+def main():
+    m_cfg = ModelConfig()
+    t_cfg = TrainConfig(batch_size=8, grad_accum_steps=4)
+    
+    tokenizer = AutoTokenizer.from_pretrained(m_cfg.encoder_name, trust_remote_code=True)
+    if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
+    
+    # 1. Load CodeXGLUE for Training
+    train_loader = prepare_data("codexglue", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="train")
+    
+    ae = SphericalAutoencoder(m_cfg).to(t_cfg.device).float()
+    # Patch pad token
+    if ae.encoder.config.pad_token_id is None:
+        ae.encoder.config.pad_token_id = tokenizer.pad_token_id
+        
+    flow = PatchedFlowDiT(m_cfg).to(t_cfg.device).float()
+    
+    trainer = Trainer(ae, flow, t_cfg, train_loader)
+    
+    # --- Training Loop ---
+    
+    # Step 1: Train AE
+    print("\n>>> Training AE on CodeXGLUE...")
+    opt_ae = optim.AdamW(filter(lambda p: p.requires_grad, ae.parameters()), lr=t_cfg.lr_ae)
+    for epoch in range(t_cfg.num_epochs_ae):
+        loss = trainer.train_ae(opt_ae)
+        print(f"AE Epoch {epoch}: Loss {loss:.4f}")
+        
+    # Step 2: Train Flow
+    print("\n>>> Training Flow Matching on CodeXGLUE...")
+    opt_flow = optim.AdamW(flow.parameters(), lr=t_cfg.lr_flow)
+    for epoch in range(t_cfg.num_epochs_flow):
+        loss = trainer.train_flow(opt_flow)
+        print(f"Flow Epoch {epoch}: Loss {loss:.4f}")
+        
+    # --- Evaluation ---
+    evaluate_on_humaneval(ae, flow, tokenizer, t_cfg.device)
+    # 训练结束后，进行 MCTS 评估
+    evaluate_with_mcts(ae, flow, tokenizer, t_cfg.device, num_samples=50)
+
+if __name__ == "__main__":
+    main()

run_wiki_flow.py

@@ -0,0 +1,282 @@
+import torch
+import torch.optim as optim
+from transformers import AutoTokenizer
+from tqdm import tqdm
+import torch.nn.functional as F
+import os
+import evaluate 
+import sacrebleu
+
+from src.config import ModelConfig, TrainConfig
+from src.models.autoencoder import ReshapedAutoencoder
+from src.models.dit import PatchedFlowDiT
+from src.trainer import Trainer
+from src.utils.data_utils import prepare_data
+
+### 加上判断eos的函数
+def _pick_stop_id(tokenizer):
+    # BERT/Jina 系通常 eos_token_id=None，用 sep_token_id 作为终止符
+    return tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.sep_token_id
+
+def _first_pos(x_1d, token_id, default):
+    # x_1d: [L]
+    idx = (x_1d == token_id).nonzero(as_tuple=True)[0]
+    return idx[0].item() if idx.numel() > 0 else default
+
+
+def calculate_metrics(sources, predictions, references):
+    """
+    计算 SARI, BLEU, 和 压缩比
+    """
+    ## 这里尝试去huggingface hub 去下载 BLEU的评估脚本，但是因为网络问题没有找到
+    # sari_metric = evaluate.load("sari")
+    # bleu_metric = evaluate.load("bleu")
+    # SARI 需要 sources
+    # sari_score = sari_metric.compute(sources=sources, predictions=predictions, references=[[r] for r in references])
+    # # BLEU
+    # bleu_score = bleu_metric.compute(predictions=predictions, references=[[r] for r in references])
+    
+    
+    # 1. BLEU
+    # sacrebleu 期望 references 是 List[List[str]] (多个参考)
+    # 这里的 references 是 List[str] (单个参考)
+    # 所以需要 transpose 一下： [[ref1, ref2, ...]]
+    bleu = sacrebleu.corpus_bleu(predictions, [references])
+    
+    # 2. SARI
+    try:
+        # corpus_sari 返回值就是一个 SARI 对象，它的 score 属性是 float
+        sari = sacrebleu.corpus_sari(sources, predictions, [references])
+        sari_score = sari.score
+    except Exception as e:
+        print(f"SARI calculation failed: {e}")
+        sari_score = 0.0
+    
+    # 3. Compression Ratio
+    ratios = [len(p) / len(s) if len(s) > 0 else 0 for p, s in zip(predictions, sources)]
+    avg_ratio = sum(ratios) / len(ratios)
+    
+    return {
+        "SARI": sari_score, # 直接使用 float
+        "BLEU": bleu.score,
+        "Compression Ratio": avg_ratio
+    }
+
+@torch.no_grad()
+def inference_batch(ae, flow, loader, tokenizer, device, steps=10, save_path="results.txt",use_oneshot=True):
+    ae.eval()
+    flow.eval()
+
+    stop_id = _pick_stop_id(tokenizer)
+    pad_id = tokenizer.pad_token_id
+    
+    print(f"\n>>> Running Inference on {len(loader.dataset)} examples...")
+    
+    all_sources = []
+    all_targets = []
+    all_generated = []
+
+    scale = getattr(ae, "latent_scale", 10.0)
+    
+    with open(save_path, "w", encoding="utf-8") as f:
+        f.write("Source\tTarget\tGenerated\n")
+        
+        for batch in tqdm(loader, desc="Inferencing"):
+            src_ids = batch['src_ids'].to(device)
+            src_mask = batch['src_mask'].to(device)
+            tgt_ids = batch['tgt_ids'].to(device)
+            B, L = src_ids.shape
+
+            z_curr = ae.encode(src_ids, src_mask)
+            z_cond = z_curr.clone()
+
+            ## 这里分别采用 one-shot 和多布采样
+            if use_oneshot:
+                # x-pred 最稳：直接 t=0 one-shot
+                t0 = torch.zeros(B, device=device)
+                z_curr = flow(z_curr, t0, condition=z_cond).float()
+            else:
+                dt = 1.0 / steps
+                for i in range(steps):
+                    t_val = i / steps
+                    # 避免 t=1 时的除零错误 (虽不常见但要防范)
+                    if t_val >= 0.999: break
+                    t = torch.ones(z_curr.shape[0], device=device) * t_val
+                    
+                    ## from v to z
+                    # v = flow(z_curr, t, condition=z_cond).float()
+                    # z_curr = z_curr + v * dt
+                    ## from z to v to zcur
+                
+                    pred_z1 = flow(z_curr, t, condition=z_cond).float()
+                    ## maybe optimize: 1 - t_val -> 1
+                    v = (pred_z1 - z_curr) / (1.0 - t_val + + 1e-4) # add epilson
+                    z_curr = z_curr + v * dt
+                    z_curr = F.normalize(z_curr, p=2, dim=-1) * scale
+                z_curr = pred_z1 # 最后一次终点预测直接使用
+            
+            z_curr = torch.nn.functional.normalize(z_curr, p=2, dim=-1) * scale ## scaling 对齐
+            
+            # ---- 3) two-pass decode to determine length by EOS ----
+            full_mask = torch.ones(B, L, device=device)  # 允许增长：全长都“可生成”
+            
+            # Pass-1: decode with full mask
+            logits1 = ae.decode(z_curr, attention_mask=full_mask)
+            ids1 = logits1.argmax(dim=-1)  # [B, L]
+
+            # find stop positions and build gen_mask
+            stop_pos = []
+            for i in range(B):
+                # 如果没预测到 stop，就用 L-1 当作“最大长度”
+                pos = _first_pos(ids1[i], stop_id, default=L - 1)
+                stop_pos.append(pos)
+            stop_pos = torch.tensor(stop_pos, device=device)
+
+            gen_mask = torch.zeros(B, L, device=device)
+            for i in range(B):
+                gen_mask[i, : stop_pos[i].item() + 1] = 1.0
+
+                # Pass-2: decode again with gen_mask, reducing tail interference
+            logits2 = ae.decode(z_curr, attention_mask=gen_mask)
+            ids2 = logits2.argmax(dim=-1)
+
+            # enforce pad after stop for clean decoding
+            ids2 = ids2.masked_fill(gen_mask == 0, pad_id)
+
+            # ---- 4) decode to text with truncation ----
+            src_texts = tokenizer.batch_decode(src_ids, skip_special_tokens=True)
+            tgt_texts = tokenizer.batch_decode(tgt_ids, skip_special_tokens=True)
+
+            gen_texts = []
+            for i in range(B):
+                end = stop_pos[i].item() + 1
+                ids_cut = ids2[i, :end]
+                gen_texts.append(tokenizer.decode(ids_cut, skip_special_tokens=True))
+    
+            # Save & Collect
+            for s, t, g in zip(src_texts, tgt_texts, gen_texts):
+                # 简单的后处理：去掉换行符以便存成 TSV
+                s_clean = s.replace("\n", " ")
+                t_clean = t.replace("\n", " ")
+                g_clean = g.replace("\n", " ")
+                
+                f.write(f"{s_clean}\t{t_clean}\t{g_clean}\n")
+                
+                all_sources.append(s_clean)
+                all_targets.append(t_clean)
+                all_generated.append(g_clean)
+                
+    return all_sources, all_targets, all_generated
+
+### add saving ckpts
+def main():
+
+    ckpt_dir = "checkpoints"
+    os.makedirs(ckpt_dir, exist_ok=True)
+    print(f"Checkpoints will be saved to: {ckpt_dir}")
+
+    # Config
+    m_cfg = ModelConfig(
+        encoder_name='../jina-embeddings-v2-base-code',
+        latent_dim=512, 
+        max_seq_len=128 # Wiki 任务文本短，用 128 足够且快
+    )
+    
+    t_cfg = TrainConfig(
+        batch_size=16, # 推理时可以大一点
+        num_epochs_ae=20,   # 增加一点 AE 训练
+        num_epochs_flow=35, # 增加 Flow 训练
+        grad_accum_steps=4,
+        use_amp=False
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(m_cfg.encoder_name, trust_remote_code=True)
+    
+    # 1. Load Data (Train & Test)
+    train_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="train")
+    test_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="test")
+    
+    # Init
+    ae = ReshapedAutoencoder(m_cfg).to(t_cfg.device).float()
+    flow = PatchedFlowDiT(m_cfg).to(t_cfg.device).float()
+    
+    if ae.encoder.config.pad_token_id is None:
+        ae.encoder.config.pad_token_id = tokenizer.pad_token_id
+        
+    # trainer = Trainer(ae, flow, t_cfg, train_loader)
+    ## 加上pad_id 和 stop_id
+    trainer = Trainer(ae, flow, t_cfg, train_loader, pad_id=tokenizer.pad_token_id, stop_id=_pick_stop_id(tokenizer))
+
+    
+    # 2. Train AE
+    opt_ae = optim.AdamW(filter(lambda p: p.requires_grad, ae.parameters()), lr=t_cfg.lr_ae)
+    best_ae_loss = float('inf')
+    print("\n>>> Start Training Autoencoder...")
+    for epoch in range(t_cfg.num_epochs_ae):
+        loss = trainer.train_ae(opt_ae)
+        print(f"AE Epoch {epoch}: Loss {loss:.4f}")
+        
+        # 保存 Best
+        if loss < best_ae_loss:
+            best_ae_loss = loss
+            torch.save(ae.state_dict(), os.path.join(ckpt_dir, "ae_best.pt"))
+            # print(f"  Saved Best AE (Loss {loss:.4f})")
+        
+        # 保存 Last (每个 epoch 覆盖，用于断点续训或检查)
+        torch.save(ae.state_dict(), os.path.join(ckpt_dir, "ae_last.pt"))
+    
+    print(f"AE Training Done. Best Loss: {best_ae_loss:.4f}")
+
+    # 3. Train Flow
+    opt_flow = optim.AdamW(flow.parameters(), lr=t_cfg.lr_flow)
+    best_flow_loss = float('inf')
+    print("\n>>> Start Training Flow DiT...")
+    for epoch in range(t_cfg.num_epochs_flow):
+        loss = trainer.train_flow(opt_flow)
+        print(f"Flow Epoch {epoch}: Loss {loss:.4f}")
+
+        # 保存 Best
+        if loss < best_flow_loss:
+            best_flow_loss = loss
+            torch.save(flow.state_dict(), os.path.join(ckpt_dir, "flow_best.pt"))
+            # print(f"  Saved Best Flow (Loss {loss:.4f})")
+            
+        # 保存 Last
+        torch.save(flow.state_dict(), os.path.join(ckpt_dir, "flow_last.pt"))
+    
+    print(f"Flow Training Done. Best Loss: {best_flow_loss:.4f}")
+    
+    # 4. Evaluation
+    # 加载最佳权重
+    ae_path = os.path.join(ckpt_dir, "ae_best.pt")
+    flow_path = os.path.join(ckpt_dir, "flow_best.pt")
+    
+    if os.path.exists(ae_path):
+        ae.load_state_dict(torch.load(ae_path, map_location=t_cfg.device))
+        print("Loaded AE Best.")
+    else:
+        print("Warning: AE Best ckpt not found, using last state.")
+
+    if os.path.exists(flow_path):
+        flow.load_state_dict(torch.load(flow_path, map_location=t_cfg.device))
+        print("Loaded Flow Best.")
+    else:
+        print("Warning: Flow Best ckpt not found, using last state.")
+
+    print("\n--- Starting Inference ---")
+    sources, targets, gens = inference_batch(
+        ae, flow, test_loader, tokenizer, t_cfg.device, 
+        steps=10, 
+        save_path="wiki_results.tsv"
+    )
+    
+    # Calculate Metrics
+    metrics = calculate_metrics(sources, gens, targets)
+    print("\n=== Metrics ===")
+    for k, v in metrics.items():
+        print(f"{k}: {v:.4f}")
+    
+    print(f"\nResults saved to wiki_results.tsv")
+
+if __name__ == "__main__":
+    main()

src/config.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+import torch
+
+@dataclass
+class ModelConfig:
+    # sequence latent space config
+    encoder_name: str = "../jina-embeddings-v2-base-code" #"jinaai/jina-embeddings-v2-base-code","microsoft/codebert-base"  # or roberta-base
+    input_dim: int = 768 # Jina Base is 768
+    latent_dim: int = 768 # 保留最大语义
+    decoder_layers: int = 4 # simple NAR decoder
+    
+    # VAE Adapter config
+    max_seq_len: int = 2048 # set according to task
+    patch_size: int = 4 # patching compress rate
+    
+    # DiT setting 
+    dit_layers: int = 12
+    dit_heads: int = 8
+    dit_hidden: int = 768 # hidden width, less than latent_dim*patch_size to cut oom
+    mlp_ratio: float = 4.0
+
+    # @property
+    # def dit_hidden(self):
+    #     return self.latent_dim
+
+@dataclass
+class TrainConfig:
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    lr_ae: float = 1e-4
+    lr_flow: float = 5e-4
+    batch_size: int = 8
+    grad_accum_steps: int = 4 # 梯度积累，等效于Batch_size = 32
+    
+    num_epochs_ae: int = 20 # 先训练AE 再训练Flow
+    num_epochs_flow: int = 50 # flow 需要训练的论数要多一些
+    
+    grad_clip: float = 1.0
+    use_amp: bool = False  # 混合精度训练,Jina+AMP 容易报错
+    save_dir: str = "./checkpoints"
+    
+    def __post_init__(self):
+        import os
+        os.makedirs(self.save_dir, exist_ok=True)

src/models/autoencoder.py

@@ -0,0 +1,181 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel, AutoConfig
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel
+
+class ResidualBlock(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * 2),
+            nn.GELU(),
+            nn.Linear(dim * 2, dim)
+        )
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        return self.norm(x + self.net(x))
+
+class ResidualAutoencoder(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        
+        # 1. Encoder (Frozen)
+        print(f"Loading Encoder: {cfg.encoder_name}...")
+        self.encoder = AutoModel.from_pretrained(cfg.encoder_name, trust_remote_code=True)
+        self.hidden_dim = self.encoder.config.hidden_size
+        for p in self.encoder.parameters(): p.requires_grad = False
+            
+        # 2. Latent Processor (No Dimension Reduction)
+        # 保持 768 维度，只做特征整理
+        # 使用残差块保证梯度流
+        self.compressor = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.hidden_dim), # 可选
+            ResidualBlock(self.hidden_dim),
+            # ResidualBlock(self.hidden_dim) # 可选：加深
+        )
+        
+        self.decompressor = nn.Sequential(
+            ResidualBlock(self.hidden_dim),
+            # ResidualBlock(self.hidden_dim),
+            nn.Linear(self.hidden_dim, self.hidden_dim)
+        )
+
+        # 3. Decoder (Pretrained)
+        print(f"Loading Decoder: {cfg.encoder_name}...")
+        self.decoder = AutoModel.from_pretrained(cfg.encoder_name, trust_remote_code=True)
+        self.decoder.config.is_decoder = False 
+        
+        # 4. Head
+        self.lm_head = nn.Linear(self.hidden_dim, self.encoder.config.vocab_size, bias=False)
+        with torch.no_grad():
+            self.lm_head.weight.copy_(self.encoder.embeddings.word_embeddings.weight)
+        self.lm_head.weight.requires_grad = True
+
+    def encode(self, input_ids, attention_mask):
+        with torch.no_grad():
+            enc_out = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
+        return self.compressor(enc_out)
+
+    def decode(self, z, attention_mask):
+        h = self.decompressor(z)
+        dec_out = self.decoder(inputs_embeds=h, attention_mask=attention_mask).last_hidden_state
+        return self.lm_head(dec_out)
+
+    def forward(self, input_ids, attention_mask):
+        z = self.encode(input_ids, attention_mask)
+        logits = self.decode(z, attention_mask)
+        return logits, z
+
+class ReshapedAutoencoder(nn.Module):
+    """
+    Sequence-to-Sequence Autoencoder with Spherical Latent Space.
+    Logic: Token -> Jina -> Linear -> Linear -> Decoder -> Token
+    """
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+
+        self.latent_scale = getattr(cfg,"latent_scale",10.0)
+        
+        # 1. Encoder (Frozen Jina)
+        print(f"Loading Pretrained Encoder: {cfg.encoder_name}...")
+        # self.encoder = AutoModel.from_pretrained(cfg.encoder_name,local_files_only=True, trust_remote_code=False)
+        self.encoder = AutoModel.from_pretrained(cfg.encoder_name,trust_remote_code=True)
+        self.hidden_dim = self.encoder.config.hidden_size
+        self.vocab_size = self.encoder.config.vocab_size
+        
+        # 冻结 Encoder 参数
+        for param in self.encoder.parameters():
+            param.requires_grad = False
+          
+        # 放弃强制 Normalize，使用 LayerNorm 进行“软约束”
+        # 结构: Hidden -> Project -> LayerNorm -> Latent
+        self.compress = nn.Sequential(
+            nn.Linear(self.hidden_dim, cfg.latent_dim),
+            nn.GELU(),
+            nn.Linear(cfg.latent_dim, cfg.latent_dim),
+            nn.LayerNorm(cfg.latent_dim) # 关键：让 latent 保持稳定分布，利于 Flow
+        )
+        
+        # 3. Decompressor
+        self.decompress = nn.Sequential(
+            nn.Linear(cfg.latent_dim, self.hidden_dim),
+            nn.GELU(),
+            nn.Linear(self.hidden_dim, self.hidden_dim),
+            nn.LayerNorm(self.hidden_dim)
+        )
+
+        # 4. Decoder (Pretrained!)
+        # <--- load from pretaining Config --->
+        print(f"Loading Pretrained Decoder: {cfg.encoder_name}...")
+        # self.decoder = AutoModel.from_pretrained(cfg.encoder_name,local_files_only=True,trust_remote_code=False)
+        self.decoder = AutoModel.from_pretrained(cfg.encoder_name,trust_remote_code=True)
+        
+        
+        # for BERT,is_decoder=False 双向 Attention，这正是 NAR 需要的
+        # 不需要 causal mask
+        self.decoder.config.is_decoder = False 
+        
+        # 5. Output Head (Trainable)
+        # 初始化为 Encoder 的 Embedding，但允许训练
+        self.lm_head = nn.Linear(self.hidden_dim, self.encoder.config.vocab_size, bias=False)
+        with torch.no_grad():
+            self.lm_head.weight.copy_(self.encoder.embeddings.word_embeddings.weight)
+        # 允许微调，以适应 decoder 输出的偏差
+        self.lm_head.weight.requires_grad = True 
+
+    def encode(self, input_ids, attention_mask):
+        """
+        Input: [B, L]
+        Output: [B, L, Latent_Dim] 
+        """
+        with torch.no_grad():
+            outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
+        
+        # Compression
+        z = self.compress(outputs.last_hidden_state) # [B, L, 768]
+
+        ## increase the scale
+        # z = z * self.latent_scale
+        
+        return z
+
+    ## 需要传入attention-mask 但是这里的疑问是对于推理没有mask 怎么办，看上去也没有判断eos
+    def decode(self, latents,attention_mask=None):
+        """
+        Input: [B, L, Latent_Dim]
+        Output: [B, L, Vocab]
+        """
+        ## back to the original scale
+        # latents = latents / self.latent_scale
+        # 1. Decompress (back to Hidden Size)
+        hidden = self.decompress(latents)
+        
+        # 2. Backbone Forward (通过 inputs_embeds 注入)
+        # AutoModel 会自动处理 mask (NAR 模式下通常是全向注意力)
+        decoder_outputs = self.decoder(
+            inputs_embeds=hidden,
+            attention_mask=attention_mask
+        )
+
+        sequence_output = decoder_outputs.last_hidden_state
+
+        # 3. Logits
+        return self.lm_head(sequence_output)
+
+    # def forward(self, input_ids, attention_mask):
+    #     z = self.encode(input_ids, attention_mask)
+    #     logits= self.decode(z, attention_mask=attention_mask)
+    #     return logits, z
+    def forward(self, input_ids, encoder_mask, decoder_mask=None):
+        if decoder_mask is None:
+            decoder_mask = encoder_mask
+        z = self.encode(input_ids, encoder_mask)
+        logits = self.decode(z, attention_mask=decoder_mask)
+        return logits, z

src/models/dit.py

@@ -0,0 +1,279 @@
+import torch
+import torch.nn as nn
+import math
+import numpy as np
+import torch.nn.functional as F
+# from timm.models.vision_transformer import Attention, Mlp -> handson_tims
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # attn = (q @ k.transpose(-2, -1)) * self.scale
+        # attn = attn.softmax(dim=-1)
+        # attn = self.attn_drop(attn)
+
+        # x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        # x = self.proj(x)
+        # x = self.proj_drop(x)
+
+        ## Replace: use Flash-Attention
+        x = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) 
+        
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class Patch1D(nn.Module):
+    """
+    [B, L, D] -> [B, L/P, D*P]
+    """
+    def __init__(self, patch_size):
+        super().__init__()
+        self.patch_size = patch_size
+
+    def forward(self, x): 
+        B, L, D = x.shape
+        # Pad sequence if not divisible by patch_size
+        # [B,31,4]->patch_size = 2 -> [B,16,8],pad is [x_31, padding_0,,,]
+        if L % self.patch_size != 0:
+            pad = self.patch_size - (L % self.patch_size)
+            x = F.pad(x, (0, 0, 0, pad))
+            
+        B, L_new, D = x.shape
+        # View as patches
+        return x.view(B, L_new // self.patch_size, D * self.patch_size)
+
+class Unpatch1D(nn.Module):
+    """
+    [B, L/P, D*P] -> [B, L, D]
+    """
+    def __init__(self, patch_size):
+        super().__init__()
+        self.patch_size = patch_size
+
+    def forward(self, x):
+        B, L_new, DP = x.shape
+        return x.view(B, L_new * self.patch_size, DP // self.patch_size)
+
+### 这里DiT的pos_embed没有使用到三角函数；另外，没有forward_with_cfg的函数实现 -> 暂时没有label_embedding
+## from: https://github.com/willisma/SiT/blob/main/models.py
+class TimestepEmbedder(nn.Module):
+    """Sinusoidal Time Embeddings"""
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size,bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size,bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        ## 兼容更多的 t 格式
+        if t.ndim > 1:
+            t = t.view(-1)
+
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+## DiTBlock, adaptive layer norm conditioning 
+class DiTBlock(nn.Module):
+    """Transformer Block with Adaptive Layer Norm (adaLN)"""
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        # c shape: [B, hidden_size]
+        # adaLN_out shape 应该是 [B, 6 * hidden_size]
+        adaLN_out = self.adaLN_modulation(c)
+        
+        # --- Debug 探针 (如果再次报错，请查看这里打印的形状) ---
+        if adaLN_out.shape[1] != 6 * self.hidden_size:
+            print(f"⚠️ DiTBlock Shape Error!")
+            print(f"Input c shape: {c.shape}")
+            print(f"adaLN output shape: {adaLN_out.shape}")
+            print(f"Expected dim1: {6 * self.hidden_size}")
+            raise ValueError("adaLN output dimension mismatch!")
+        # ----------------------------------------------------
+
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaLN_out.chunk(6, dim=1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+class PatchedFlowDiT(nn.Module):
+    """
+    Main DiT Architecture for Flow Matching
+    Input: z_t (Noisy Latent) + t (Time) + condition (Original Latent)
+    Output: velocity vector
+    """
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        ## add patch and unpatch block here
+        self.patcher = Patch1D(cfg.patch_size)
+        self.unpatcher = Unpatch1D(cfg.patch_size)
+
+        # 计算 Patch 后的输入维度
+        # Input to DiT = Patch(z_t) + Patch(Condition)
+        # 维度 = (Latent * Patch) * 2
+        input_feat_dim = cfg.latent_dim * cfg.patch_size
+        
+        # Projection to DiT Hidden Size
+        self.input_proj = nn.Linear(input_feat_dim * 2, cfg.dit_hidden)
+
+        # Time & Pos Embeddings
+        self.time_embed = TimestepEmbedder(cfg.dit_hidden)
+        patched_len = (cfg.max_seq_len + cfg.patch_size - 1) // cfg.patch_size
+        self.pos_embed = nn.Parameter(torch.zeros(1, patched_len, cfg.dit_hidden))
+        
+        self.blocks = nn.ModuleList([
+            DiTBlock(cfg.dit_hidden, cfg.dit_heads) for _ in range(cfg.dit_layers)
+        ])
+        
+        # Output Projection (Predict Velocity)
+        self.final_layer = nn.Linear(cfg.dit_hidden, input_feat_dim)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize pos_embed
+        nn.init.normal_(self.pos_embed, std=0.02)
+
+        # Zero-out adaLN modulation layers
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out final layer -> modify: to predict data, so initialize is xavier or normal
+        # nn.init.constant_(self.final_layer.weight, 0)
+        # nn.init.constant_(self.final_layer.bias, 0)
+        nn.init.xavier_uniform_(self.final_layer.weight)
+        nn.init.constant_(self.final_layer.bias, 0)
+
+    def forward(self, z_t, t, condition):
+        # x: [Batch, Seq, Dim]
+        # t: [Batch]
+        # condition: [Batch, Seq, Dim] (Optional, e.g., Source Sentence)
+        """
+        z_t: [B, L, D]
+        condition: [B, L, D]
+        """
+        # 1. Patching
+        z_p = self.patcher(z_t)
+        c_p = self.patcher(condition)
+        
+        # 2. Concat & Project(Jit Style: Explicit Conditioning)
+        x = torch.cat([z_p, c_p], dim=-1)
+        x = self.input_proj(x)
+        
+        # 3. Add Embeddings
+        t_emb = self.time_embed(t)
+        # Handle length mismatch due to padding
+        L_curr = x.shape[1]
+        x = x + self.pos_embed[:, :L_curr, :]
+        
+        # 4. Transformer
+        for block in self.blocks:
+            x = block(x, t_emb)
+            
+        # 5. Output & Unpatch
+        v_p = self.final_layer(x)
+        v = self.unpatcher(v_p)
+        
+        # Crop to original length
+        return v[:, :z_t.shape[1], :]
+
+    def forward_with_cfg(self, x, t, condition, cfg_scale):
+        """
+        支持 Classifier-Free Guidance 的前向传播
+        """
+        # 1. condition
+        cond_out = self.forward(x, t, condition)
+        
+        # 2.uncondition
+        uncond_out = self.forward(x, t, condition=None)
+        
+        # 3. classifier-free guidance
+        # eps = eps_uncond + s * (eps_cond - eps_uncond)
+        return uncond_out + cfg_scale * (cond_out - uncond_out)

src/search.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+class DiffuMCTS:
+    """
+    Diffusion-based Monte Carlo Tree Search (Simulation-based).
+    Flow Matching Generation -> Rollout.
+    Sandbox to evaluation.
+    """
+    def __init__(self, ae, flow, tokenizer, sandbox, device, config):
+        self.ae = ae
+        self.flow = flow
+        self.tokenizer = tokenizer
+        self.sandbox = sandbox
+        self.device = device
+        
+        # 搜索配置
+        self.num_branches = 8      # 分支数量 (K)
+        self.split_t = 0.5         # 在哪个时间点分叉 (0=Target, 1=Source)
+        self.noise_scale = 0.1     # 分支时的扰动强度
+        self.steps = 10            # Flow ODE 步数
+        
+    @torch.no_grad()
+    def solve(self, buggy_code, test_code, entry_point):
+        """
+        对外接口：尝试修复代码
+        Returns:
+            fixed_code (str): 修复后的代码
+            success (bool): 是否通过测试
+        """
+        # 1. Encode Buggy Code (Source) -> z_1
+        tokens = self.tokenizer(
+            buggy_code, 
+            max_length=2048, 
+            padding="max_length", 
+            truncation=True, 
+            return_tensors="pt"
+        ).to(self.device)
+        
+        z_buggy = self.ae.encode(tokens['input_ids'], tokens['attention_mask'])
+        
+        # 2. Search Strategy: Parallel Branching
+        # 我们执行一次带有分支的推理
+        best_code, success = self._parallel_rollout(z_buggy, test_code, entry_point)
+        
+        return best_code, success
+
+    def _parallel_rollout(self, z_start, test_code, entry_point):
+        """
+        执行并行 Rollout 搜索
+        """
+        B, L, D = z_start.shape
+        K = self.num_branches
+        
+        # --- Stage 1: Deterministic Flow (1.0 -> split_t) ---
+        # 先从 Buggy 状态走几步，让语义稍微稳定一点
+        z_curr = z_start.clone()
+        z_cond = z_start.clone() # Condition 始终是 Buggy Code
+        
+        dt = 1.0 / self.steps
+        # 计算 split 对应的步数索引
+        split_step_idx = int((1.0 - self.split_t) * self.steps)
+        
+        # 走前半程
+        for i in range(split_step_idx):
+            t_val = 1.0 - (i / self.steps) # 从 1 走向 0
+            t_tensor = torch.ones(B, device=self.device) * t_val
+            
+            # ODE Step: z_next = z_prev - v * dt (注意时间方向)
+            # Rectified Flow 定义 v = z_1 - z_0 (从 Good 到 Bad 的反向? 或者是 Bad 到 Good?)
+            # 我们训练时: z_t = (1-t)z_bad + t*z_good. 
+            # 所以 t=0 是 Bad, t=1 是 Good.
+            # 为了方便，我们定义 t 从 0 (Bad) 走向 1 (Good)。
+            
+            # 修正逻辑：
+            # Forward Euler: z_{t+dt} = z_t + v * dt
+            # t 从 0 增加到 split_t
+            current_t_val = i / self.steps
+            t_tensor = torch.ones(B, device=self.device) * current_t_val
+            
+            v = self.flow(z_curr, t_tensor, condition=z_cond)
+            z_curr = z_curr + v * dt
+
+        # --- Stage 2: Expansion (Branching) ---
+        # 复制 K 份
+        # [B, L, D] -> [B*K, L, D]
+        z_branches = z_curr.repeat(K, 1, 1)
+        z_cond_branches = z_cond.repeat(K, 1, 1)
+        
+        # 注入高斯噪声 (Exploration)
+        # z' = z + noise
+        noise = torch.randn_like(z_branches) * self.noise_scale
+        z_branches = z_branches + noise
+        
+        # 重新投影回球面 (保持流形约束)
+        z_branches = F.normalize(z_branches, p=2, dim=-1)
+
+        # --- Stage 3: Rollout (split_t -> 1.0) ---
+        # 并行推演所有分支
+        remaining_steps = self.steps - split_step_idx
+        
+        for i in range(remaining_steps):
+            step_idx = split_step_idx + i
+            current_t_val = step_idx / self.steps
+            
+            # [B*K]
+            t_tensor = torch.ones(z_branches.shape[0], device=self.device) * current_t_val
+            
+            v = self.flow(z_branches, t_tensor, condition=z_cond_branches)
+            z_branches = z_branches + v * dt
+
+        # --- Stage 4: Decoding & Verification ---
+        # 批量解码
+        # [B*K, L, D] -> [B*K, L, Vocab]
+        logits = self.ae.decode(z_branches)
+        pred_ids = torch.argmax(logits, dim=-1)
+        
+        candidate_codes = self.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+        
+        # 验证 Loop
+        # 只要有一个通过，就算成功 (Pass@k)
+        for code in candidate_codes:
+            is_pass, msg = self.sandbox.run(code, test_code, entry_point)
+            if is_pass:
+                return code, True
+        
+        # 如果都失败，返回第一个（或者可以设计 heuristic 选择最接近的）
+        return candidate_codes[0], False

src/trainer.py

@@ -0,0 +1,267 @@
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+class Trainer:
+    def __init__(self, ae, flow, cfg, loader, pad_id, stop_id):
+        self.ae = ae.to(cfg.device)
+        self.flow = flow.to(cfg.device) if flow else None
+        self.cfg = cfg
+        self.loader = loader
+        self.device = cfg.device
+        self.pad_id = pad_id
+        self.stop_id = stop_id
+
+    def train_ae(self, optimizer):
+        self.ae.train()
+        total_loss = 0
+        pbar = tqdm(self.loader, desc="Train AE")
+        optimizer.zero_grad()
+        
+        for step, batch in enumerate(pbar):
+            tgt = batch['tgt_ids'].to(self.device)
+            mask = batch['tgt_mask'].to(self.device)
+            
+            # logits, z = self.ae(tgt, mask)
+
+            # ## 不太明白这里的mask 机制
+            # labels = tgt.masked_fill(mask == 0, -100)
+            # loss = F.cross_entropy(
+            #     logits.view(-1, logits.size(-1)),
+            #     labels.view(-1),
+            #     ignore_index=-100
+            # )
+            # Reconstruction Loss
+            # loss = F.cross_entropy(logits.view(-1, logits.size(-1)), tgt.view(-1), ignore_index=1)
+            logits, z = self.ae(tgt, mask)   # decoder_mask 默认 = mask
+
+            V = logits.size(-1)
+            B, L = tgt.shape3
+
+            # 1) token loss：只看 mask==1
+            labels_tok = tgt.masked_fill(mask == 0, -100)
+            loss_tok = F.cross_entropy(
+                logits.view(-1, V),
+                labels_tok.view(-1),
+                ignore_index=-100,
+                reduction="mean"
+            )
+
+            # 2) pad loss：mask==0 的位置强制预测 PAD（轻权重）
+            pad_pos = (mask == 0)
+            if pad_pos.any():
+                # 每个位置的 CE
+                ce_all = F.cross_entropy(
+                    logits.view(-1, V),
+                    tgt.new_full((B * L,), self.pad_id),
+                    reduction="none"
+                ).view(B, L)
+                loss_pad = (ce_all * pad_pos).sum() / (pad_pos.sum() + 1e-6)
+            else:
+                loss_pad = logits.new_tensor(0.0)
+
+            # 3) 可选：stop 位置加权（让 SEP 更稳）
+            stop_pos = ((tgt == self.stop_id) & (mask == 1))
+            if stop_pos.any():
+                ce_tok = F.cross_entropy(
+                    logits.view(-1, V),
+                    tgt.view(-1),
+                    reduction="none"
+                ).view(B, L)
+                loss_stop = (ce_tok * stop_pos).sum() / (stop_pos.sum() + 1e-6)
+            else:
+                loss_stop = logits.new_tensor(0.0)
+
+            # 合成：pad/stop 的权重别太大
+            lambda_pad = 0.1
+            lambda_stop = 0.2
+            loss = loss_tok + lambda_pad * loss_pad + lambda_stop * loss_stop
+
+            loss = loss / self.cfg.grad_accum_steps
+            loss.backward()
+            
+            if (step + 1) % self.cfg.grad_accum_steps == 0:
+                optimizer.step()
+                optimizer.zero_grad()
+                
+            total_loss += loss.item() * self.cfg.grad_accum_steps
+            pbar.set_postfix(loss=loss.item() * self.cfg.grad_accum_steps)
+            
+        return total_loss / len(self.loader)
+
+    def train_robust_ae(self, optimizer):
+        
+        self.ae.train()
+        total_loss = 0
+        noise_std = 0.05 
+
+        for batch in self.loader:
+            tgt_ids = batch['tgt_ids'].to(self.device)
+            tgt_mask = batch['tgt_mask'].to(self.device)
+            
+            # 1. get normal z
+            with torch.no_grad():
+                z_clean = self.ae.encode(tgt_ids, tgt_mask)
+            
+            # 2. add noise (Denoising Training)
+            #  Decoder -> like z 
+            noise = torch.randn_like(z_clean) * noise_std
+            z_noisy = z_clean + noise
+            
+            # 3. Decode
+            logits = self.ae.decode(z_noisy, attention_mask=tgt_mask)
+
+            # 4. Loss
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), 
+                                   tgt_ids.view(-1), 
+                                   reduction='none')
+            loss = (loss * tgt_mask.view(-1)).sum() / tgt_mask.sum()
+
+            # Backward
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            
+            total_loss += loss.item()
+            
+        return total_loss / len(self.loader)
+
+
+    def train_ae_combined(self, optimizer, epoch, max_epochs):
+        """
+        结合了 基础重建 + 课程去噪 + Pad/Stop 优化
+        """
+        self.ae.train()
+        total_loss = 0
+        
+        # --- 课程噪声调度 (Curriculum Noise) ---
+        # 前 20% 的 Epoch 不加噪声，先学好重建
+        # 后面线性增加到 0.1
+        if epoch < max_epochs * 0.2:
+            current_noise = 0.0
+        else:
+            progress = (epoch - max_epochs * 0.2) / (max_epochs * 0.8)
+            current_noise = 0.1 * progress # 最大噪声 0.1
+
+        pbar = tqdm(self.loader, desc=f"Train AE (Noise={current_noise:.4f})")
+        
+        for step, batch in enumerate(pbar):
+            tgt = batch['tgt_ids'].to(self.device)
+            mask = batch['tgt_mask'].to(self.device)
+            
+            # 1. Encode Clean
+            with torch.no_grad():
+                z_clean = self.ae.encode(tgt, mask)
+            
+            # 2. Add Noise (如果 noise > 0)
+            if current_noise > 0:
+                noise = torch.randn_like(z_clean) * current_noise
+                z_input = z_clean + noise
+            else:
+                z_input = z_clean
+            
+            # 3. Decode
+            logits = self.ae.decode(z_input, attention_mask=mask)
+            
+            # 4. Calculate Advanced Loss (Copy from your original code)
+            V = logits.size(-1)
+            B, L = tgt.shape
+
+            # Token Loss (只看 mask==1)
+            labels_tok = tgt.masked_fill(mask == 0, -100)
+            loss_tok = F.cross_entropy(logits.view(-1, V), labels_tok.view(-1), ignore_index=-100)
+
+            # Pad Loss (mask==0)
+            pad_pos = (mask == 0)
+            if pad_pos.any():
+                ce_pad = F.cross_entropy(logits.view(-1, V), tgt.new_full((B*L,), self.pad_id), reduction='none').view(B,L)
+                loss_pad = (ce_pad * pad_pos).sum() / (pad_pos.sum() + 1e-6)
+            else:
+                loss_pad = torch.tensor(0.0, device=self.device)
+                
+            # Stop Loss
+            stop_pos = ((tgt == self.stop_id) & (mask == 1))
+            if stop_pos.any():
+                ce_stop = F.cross_entropy(logits.view(-1, V), tgt.view(-1), reduction='none').view(B,L)
+                loss_stop = (ce_stop * stop_pos).sum() / (stop_pos.sum() + 1e-6)
+            else:
+                loss_stop = torch.tensor(0.0, device=self.device)
+
+            # 合并 Loss
+            loss = loss_tok + 0.1 * loss_pad + 0.5 * loss_stop # 提高一点 stop 的权重
+            
+            # Backward
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            
+            total_loss += loss.item()
+            pbar.set_postfix(loss=loss.item())
+            
+        return total_loss / len(self.loader)
+
+    def train_flow(self, optimizer):
+        self.flow.train()
+        self.ae.eval()
+        total_loss = 0
+        pbar = tqdm(self.loader, desc="Train Flow")
+        optimizer.zero_grad()
+
+        scale = getattr(self.ae, "latent_scale", 10.0)
+        
+        for step, batch in enumerate(pbar):
+            src = batch['src_ids'].to(self.device)
+            src_mask = batch['src_mask'].to(self.device)
+            tgt = batch['tgt_ids'].to(self.device)
+            tgt_mask = batch['tgt_mask'].to(self.device)
+            
+            with torch.no_grad():
+                z_bad = self.ae.encode(src, src_mask) # norm ~ scale
+                z_good = self.ae.encode(tgt, tgt_mask) # norm ~ scale
+                
+            # Rectified Flow
+            bs = z_bad.shape[0]
+            t = torch.rand(bs, device=self.device).view(bs, 1, 1)
+            
+            
+            # Interpolation: Bad -> Good, modify-> push back to sphere
+            z_t_linear = (1 - t) * z_bad + t * z_good
+            ## test before or after
+            # z_t = F.normalize(z_t_linear, p=2, dim=-1) * scale
+            z_t = z_t_linear
+
+            # Modify: pred_v to pred_x
+            # target_v = z_good - z_bad
+            # pred_v = self.flow(z_t, t.squeeze(), condition=z_bad)
+            # loss = F.mse_loss(pred_v, target_v)
+            
+            # to predict z_good (Target)
+            pred_z1 = self.flow(z_t, t, condition=z_bad)
+            # 3) (强烈建议) 把输出也投影回同一球面，避免 off-manifold -> 都不要normalize
+            pred_z1 = pred_z1
+            # pred_z1 = F.normalize(pred_z1, p=2, dim=-1) * scale
+            # Loss 直接算与 z_good 的距离
+            ## 修改：loss必须按照mask 算有效token
+            mse = (pred_z1 - z_good).pow(2).mean(dim=-1)  # [B,L]
+            w = tgt_mask.float()
+
+            # stop 位置加权
+            stop_pos = ((tgt == self.stop_id) & (tgt_mask == 1))
+            w = w + stop_pos.float() * 2.0   # 让 SEP 位置权重更大（比如 +2）
+
+            loss = (mse * w).sum() / (w.sum() + 1e-6)
+
+            # loss = (mse * tgt_mask).sum() / (tgt_mask.sum() + 1e-6)
+            # loss = F.mse_loss(pred_z1, z_good)
+
+            loss = loss / self.cfg.grad_accum_steps
+            loss.backward()
+            
+            if (step + 1) % self.cfg.grad_accum_steps == 0:
+                optimizer.step()
+                optimizer.zero_grad()
+                
+            total_loss += loss.item() * self.cfg.grad_accum_steps
+            pbar.set_postfix(loss=loss.item() * self.cfg.grad_accum_steps)
+            
+        return total_loss / len(self.loader)

src/utils/data_utils.py

@@ -0,0 +1,157 @@
+import os
+from datasets import load_dataset
+from transformers import AutoTokenizer
+import torch
+from torch.utils.data import DataLoader
+
+def custom_collate(batch):
+    return {
+        'src_ids': torch.stack([torch.tensor(x['src_ids']) for x in batch]),
+        'src_mask': torch.stack([torch.tensor(x['src_mask']) for x in batch]),
+        'tgt_ids': torch.stack([torch.tensor(x['tgt_ids']) for x in batch]),
+        'tgt_mask': torch.stack([torch.tensor(x['tgt_mask']) for x in batch]),
+        # 保留测试用例用于验证 (仅 Eval 时有效)
+        'test_code': [x.get('test_code', "") for x in batch],
+        'entry_point': [x.get('entry_point', "") for x in batch]
+    }
+
+def prepare_data(task_name, tokenizer, max_len, batch_size, split="train"):
+    """
+    支持 split 参数，方便划分训练集和测试集
+    """
+    print(f"Loading {task_name} ({split})...")
+    
+    if task_name == "codexglue":
+        # 训练集：Microsoft CodeXGLUE (Python Refinement)
+        # 包含 GitHub Bug -> Fix
+        dataset = load_dataset("./code_x_glue_cc_code_refinement_full", "medium", split=split)
+
+        # 40k
+        if split == "train": dataset = dataset.select(range(40000)) 
+
+        # Case A: 标准修复数据 (有 source 和 target)
+        if 'source' in cols and 'target' in cols:
+            print(">> Detected standard refinement pairs.")
+            def preprocess_standard(ex):
+                src = tokenizer(ex['source'], max_length=max_len, padding="max_length", truncation=True)
+                tgt = tokenizer(ex['target'], max_length=max_len, padding="max_length", truncation=True)
+                return {
+                    'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
+                    'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask']
+                }
+            preprocess_fn = preprocess_standard
+
+        # Case B: 只有代码 (有 code)，需要人工注入 Bug
+        elif 'code' in cols:
+            print(">> Detected raw code. Not to inject synthetic bugs...")
+            
+        else:
+            raise ValueError(f"Dataset columns {cols} not recognized. Need 'source'/'target' or 'code'.")
+        
+        def preprocess(ex):
+            buggy = ex['source']
+            fixed = ex['target']
+            
+            src = tokenizer(buggy, max_length=max_len, padding="max_length", truncation=True)
+            tgt = tokenizer(fixed, max_length=max_len, padding="max_length", truncation=True)
+            
+            return {
+                'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
+                'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask']
+            }
+        
+        # 移除原始列
+        cols = dataset.column_names
+    
+    elif task_name == "humanevalpack":
+        # 验证集：HumanEvalPack (Fix Task)
+        # 包含 Buggy Code 和 对应的 Unit Tests
+        dataset = load_dataset("./bigcode_humanevalpack_full", "python", split="test") # 只有 test 集
+        
+        # 筛选出 FIX 任务
+        dataset = dataset.filter(lambda x: x['task_id'].startswith("Python/FIX"))
+        
+        def preprocess(ex):
+            # prompt 是前面的描述，buggy_solution 是有 bug 的代码
+            # 为了简化，我们把 prompt + buggy_solution 作为输入
+            full_buggy = ex['prompt'] + "\n" + ex['buggy_solution']
+            full_fixed = ex['prompt'] + "\n" + ex['canonical_solution']
+            
+            src = tokenizer(full_buggy, max_length=max_len, padding="max_length", truncation=True)
+            tgt = tokenizer(full_fixed, max_length=max_len, padding="max_length", truncation=True)
+            
+            return {
+                'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
+                'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask'],
+                'test_code': ex['test'],         # 核心：保留测试代码
+                'entry_point': ex['entry_point'] # 核心：保留入口函数名
+            }
+        
+        # 保留所有列用于 debug，dataset.map 会自动处理返回的 dict
+        cols = [] # 不自动删除列，我们需要 test 列在 collate 中处理
+    
+    # --- 1. Load Dataset ---
+    elif task_name == "wiki":
+         # 尝试本地加载，失败则下载
+        try:
+            dataset = load_dataset("./wikilarge-dataset")
+        except:
+            print("Local load failed, downloading from Hub...")
+            dataset = load_dataset("wikilarge")
+        
+        # 手动划分: train用前10000条, test用后1000条 (做demo够了，全量太慢)
+        if split == "train":
+            dataset = dataset['train'].select(range(20000))
+        else:
+            # 假设总共有 ~290k，我们取后面一点做测试
+            dataset = dataset['train'].select(range(20000, 25000))
+
+        # 自动探测列名
+        cols = dataset.column_names
+        print(f"Wiki Dataset Columns: {cols}")
+        
+        # 映射列名到 src/tgt
+        if 'src' in cols and 'dst' in cols:
+            src_key, tgt_key = 'src', 'dst'
+        elif 'Normal' in cols and 'Simple' in cols:
+            src_key, tgt_key = 'Normal', 'Simple'
+        else:
+            raise ValueError(f"Unknown column format for WikiLarge: {cols}")
+
+        def preprocess(ex):
+            # Source (Complex) -> Target (Simple)
+            src = tokenizer(ex[src_key], max_length=max_len, padding="max_length", truncation=True)
+            tgt = tokenizer(ex[tgt_key], max_length=max_len, padding="max_length", truncation=True)
+            return {
+                'src_ids': src['input_ids'], 'src_mask': src['attention_mask'],
+                'tgt_ids': tgt['input_ids'], 'tgt_mask': tgt['attention_mask']
+            }
+        
+    elif task_name == "mbpp":
+        dataset = load_dataset("mbpp", split="train[:500]")
+        print(f"MBPP Dataset Columns: {dataset.column_names}")
+        
+        # MBPP 自重建任务: src=code, tgt=code
+        def preprocess(ex):
+            enc = tokenizer(ex['code'], max_length=max_len, padding="max_length", truncation=True)
+            return {
+                'src_ids': enc['input_ids'], 'src_mask': enc['attention_mask'],
+                'tgt_ids': enc['input_ids'], 'tgt_mask': enc['attention_mask']
+            }
+    
+    else:
+        raise ValueError(f"Unknown task: {task_name}")
+
+    # --- 2. Map & Batch ---
+    print(f"Preprocessing {task_name} data...")
+    # 使用 remove_columns=dataset.column_names 确保删除所有原始列
+    print(f"Preprocessing {len(dataset)} examples...")
+    dataset = dataset.map(
+        preprocess, 
+        batched=True, 
+        remove_columns=dataset.column_names, 
+        num_proc=4
+    )
+    
+    # Test 集不 shuffle，方便对齐
+    return DataLoader(dataset, batch_size=batch_size, shuffle=(split=="train"), collate_fn=custom_collate)

src/utils/sandbox.py

@@ -0,0 +1,156 @@
+import sys
+import io
+import multiprocessing
+import contextlib
+import signal
+import subprocess
+import tempfile
+import shutil
+
+# 增加超时控制
+class TimeoutException(Exception): pass
+
+def timeout_handler(signum, frame):
+    raise TimeoutException
+
+def _exec_code(code_str, test_code, entry_point, result_queue):
+    """
+    运行生成的代码 + 测试用例
+    """
+    capture = io.StringIO()
+    success = False
+    error_msg = ""
+    
+    try:
+        # 简单的超时机制 (Linux only, Windows需要其他方式)
+        # signal.signal(signal.SIGALRM, timeout_handler)
+        # signal.alarm(2) # 2秒超时
+        
+        with contextlib.redirect_stdout(capture), contextlib.redirect_stderr(capture):
+            # 创建独立命名空间
+            scope = {}
+            
+            # 1. 执行生成的代码 (定义函数)
+            exec(code_str, scope)
+            
+            # 2. 检查入口函数是否存在
+            if entry_point not in scope:
+                raise ValueError(f"Entry point {entry_point} not found in generated code.")
+            
+            # 3. 执行测试用例
+            # HumanEval 的测试用例通常是 "check(entry_point_func)" 的形式
+            # 我们需要把 check 函数定义也 exec 进去，或者拼接到一起
+            full_test_script = code_str + "\n" + test_code + f"\ncheck({entry_point})"
+            
+            exec(full_test_script, scope)
+            
+            success = True
+            
+    except Exception as e:
+        error_msg = str(e)
+    # finally:
+        # signal.alarm(0)
+    
+    result_queue.put((success, error_msg))
+
+class SafeSandbox:
+    def __init__(self, timeout=5.0):
+        self.timeout = timeout
+
+    def run(self, code, test_code, entry_point):
+        queue = multiprocessing.Queue()
+        p = multiprocessing.Process(target=_exec_code, args=(code, test_code, entry_point, queue))
+        p.start()
+        p.join(self.timeout)
+
+        if p.is_alive():
+            p.terminate()
+            p.join()
+            return False, "Timeout"
+
+        if not queue.empty():
+            return queue.get()
+        return False, "Unknown Error"
+
+class JavaSandbox:
+    def __init__(self, timeout=5.0):
+        self.timeout = timeout
+        
+        # 检查 Java 环境
+        if shutil.which("javac") is None or shutil.which("java") is None:
+            raise RuntimeError("Java environment (jdk) not found. Please install java.")
+
+    def run(self, code, test_code, entry_point):
+        """
+        code: 修复后的 Java 方法代码
+        test_code: 包含 main 函数的测试类代码，调用 entry_point
+        entry_point: 方法名 (Java 中通常不需要，只要 test_code 写对)
+        """
+        # 创建临时目录
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_name = "Solution.java" # 假设类名是 Solution
+            file_path = os.path.join(temp_dir, file_name)
+            
+            # 拼接代码：我们需要把生成的 method 塞进一个 Class 里
+            # 这里假设 code 只是一个 method，test_code 是 main 函数
+            # 你需要根据数据集的具体格式调整拼接逻辑
+            
+            full_source = f"""
+public class Solution {{
+    {code}
+
+    {test_code}
+}}
+"""
+            # 1. 写入文件
+            with open(file_path, "w") as f:
+                f.write(full_source)
+            
+            # 2. 编译
+            compile_cmd = ["javac", file_path]
+            try:
+                subprocess.run(compile_cmd, check=True, capture_output=True, timeout=10)
+            except subprocess.CalledProcessError as e:
+                return False, f"Compilation Error: {e.stderr.decode()}"
+            except subprocess.TimeoutExpired:
+                return False, "Compilation Timeout"
+
+            # 3. 运行
+            run_cmd = ["java", "-cp", temp_dir, "Solution"]
+            try:
+                result = subprocess.run(run_cmd, capture_output=True, timeout=self.timeout)
+                if result.returncode == 0:
+                    return True, result.stdout.decode()
+                else:
+                    return False, f"Runtime Error: {result.stderr.decode()}"
+            except subprocess.TimeoutExpired:
+                return False, "Runtime Timeout"
+
+# 单元测试
+if __name__ == "__main__":
+    sandbox = JavaSandbox()
+    
+    # 正确代码
+    code_pass = """
+    public static int add(int a, int b) {
+        return a + b;
+    }
+    """
+    test_pass = """
+    public static void main(String[] args) {
+        if (add(1, 1) == 2) {
+            System.out.println("PASS");
+        } else {
+            System.exit(1);
+        }
+    }
+    """
+    print("Test Pass:", sandbox.run(code_pass, test_pass, "add"))
+    
+    # 错误代码
+    code_fail = """
+    public static int add(int a, int b) {
+        return a * b; // Bug
+    }
+    """
+    print("Test Fail:", sandbox.run(code_fail, test_pass, "add"))

tests/test_models.py

@@ -0,0 +1,61 @@
+import unittest
+import torch
+import sys
+import os
+
+# 确保能导入 src
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from src.config import ModelConfig
+from src.models.autoencoder import LatentAutoencoder
+from src.models.dit import FlowDiT
+
+class TestModels(unittest.TestCase):
+    def setUp(self):
+        # 构造一个小配置用于测试
+        self.cfg = ModelConfig(
+            encoder_name="roberta-base", 
+            latent_dim=128, 
+            max_seq_len=32,
+            decoder_layers=2, # 快一点
+            dit_layers=2
+        )
+        # 强行覆盖 dit_hidden 确保测试一致性 (虽然 property 已经保证了)
+        # self.cfg.dit_hidden = 128 
+
+    def test_ae_shape(self):
+        print("\nTesting Autoencoder Shape...")
+        model = LatentAutoencoder(self.cfg)
+        input_ids = torch.randint(0, 100, (2, 32))
+        mask = torch.ones((2, 32))
+        logits, z = model(input_ids, mask)
+        
+        self.assertEqual(z.shape, (2, 32, 128))
+        # 50265 是 RoBERTa 的词表大小
+        self.assertEqual(logits.shape, (2, 32, 50265)) 
+        print("AE Shape Check Passed.")
+
+    def test_dit_shape(self):
+        print("\nTesting DiT Shape...")
+        model = FlowDiT(self.cfg)
+        x = torch.randn(2, 32, 128) # B, Seq, Dim
+        t = torch.rand(2)           # B
+        cond = torch.randn(2, 32, 128)
+        
+        out = model(x, t, condition=cond)
+        self.assertEqual(out.shape, (2, 32, 128))
+        print("DiT Shape Check Passed.")
+
+    def test_cfg_forward(self):
+        print("\nTesting CFG Forward...")
+        model = FlowDiT(self.cfg)
+        x = torch.randn(2, 32, 128)
+        t = torch.rand(2)
+        cond = torch.randn(2, 32, 128)
+        
+        out = model.forward_with_cfg(x, t, cond, cfg_scale=3.0)
+        self.assertEqual(out.shape, (2, 32, 128))
+        print("CFG Check Passed.")
+
+if __name__ == "__main__":
+    unittest.main()

train_ae.py

@@ -0,0 +1,86 @@
+import torch
+import torch.optim as optim
+from transformers import AutoTokenizer
+import os
+import argparse
+
+from src.config import ModelConfig, TrainConfig
+from src.models.autoencoder import ReshapedAutoencoder,ResidualAutoencoder
+from src.trainer import Trainer
+from src.utils.data_utils import prepare_data
+
+def _pick_stop_id(tokenizer):
+    return tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.sep_token_id
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--save_dir", type=str, default="/mnt/hdfs/user/lixinyu.222/CodeFlow/robust_checkpoints", help="Directory to save checkpoints")
+    args = parser.parse_args()
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    print(f"Checkpoints will be saved to: {args.save_dir}")
+
+    # --- Config ---
+    m_cfg = ModelConfig(
+        encoder_name='../jina-embeddings-v2-base-code', # 请根据实际路径修改
+        latent_dim=512, 
+        max_seq_len=128
+    )
+    
+    t_cfg = TrainConfig(
+        batch_size=16,
+        num_epochs_ae=20, # 只关注 AE 的 epoch
+        grad_accum_steps=4,
+        use_amp=False,
+        lr_ae=1e-4
+    )
+    
+    # --- Data & Tokenizer ---
+    tokenizer = AutoTokenizer.from_pretrained(m_cfg.encoder_name, local_files_only=True, trust_remote_code=False)
+    train_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="train")
+    
+    # --- Model ---
+    ae = ReshapedAutoencoder(m_cfg).to(t_cfg.device).float()
+    # ae = ResidualAutoencoder(m_cfg).to(t_cfg.device).float()
+    
+    if ae.encoder.config.pad_token_id is None:
+        ae.encoder.config.pad_token_id = tokenizer.pad_token_id
+
+    # --- Trainer ---
+    # 这里 flow 传 None，因为只训 AE
+    trainer = Trainer(
+        ae=ae, 
+        flow=None, 
+        cfg=t_cfg, 
+        loader=train_loader, 
+        pad_id=tokenizer.pad_token_id, 
+        stop_id=_pick_stop_id(tokenizer)
+    )
+    
+    # --- Optimizer ---
+    opt_ae = optim.AdamW(filter(lambda p: p.requires_grad, ae.parameters()), lr=t_cfg.lr_ae)
+    
+    # --- Training Loop ---
+    best_ae_loss = float('inf')
+    print("\n>>> Start Training Autoencoder...")
+    
+    for epoch in range(t_cfg.num_epochs_ae):
+        # loss = trainer.train_ae(opt_ae)
+        # loss = trainer.train_robust_ae(opt_ae)
+        loss = trainer.train_ae_combined(opt_ae, epoch, t_cfg.num_epochs_ae)
+        print(f"AE Epoch {epoch}: Loss {loss:.4f}")
+        
+        # Save Best
+        if loss < best_ae_loss:
+            best_ae_loss = loss
+            save_path = os.path.join(args.save_dir, "ae_best.pt")
+            torch.save(ae.state_dict(), save_path)
+            print(f"  Saved Best AE to {save_path}")
+        
+        # Save Last
+        torch.save(ae.state_dict(), os.path.join(args.save_dir, "ae_last.pt"))
+    
+    print(f"AE Training Done. Best Loss: {best_ae_loss:.4f}")
+
+if __name__ == "__main__":
+    main()

train_flow.py

@@ -0,0 +1,227 @@
+import torch
+import torch.optim as optim
+from transformers import AutoTokenizer
+from tqdm import tqdm
+import torch.nn.functional as F
+import os
+import argparse
+import sacrebleu
+
+from src.config import ModelConfig, TrainConfig
+from src.models.autoencoder import ReshapedAutoencoder
+from src.models.dit import PatchedFlowDiT
+from src.trainer import Trainer
+from src.utils.data_utils import prepare_data
+
+# --- Helper Functions for Inference (复制过来以便独立运行) ---
+def _pick_stop_id(tokenizer):
+    return tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.sep_token_id
+
+def _first_pos(x_1d, token_id, default):
+    idx = (x_1d == token_id).nonzero(as_tuple=True)[0]
+    return idx[0].item() if idx.numel() > 0 else default
+
+def calculate_metrics(sources, predictions, references):
+    bleu = sacrebleu.corpus_bleu(predictions, [references])
+    try:
+        sari = sacrebleu.corpus_sari(sources, predictions, [references])
+        sari_score = sari.score
+    except Exception:
+        sari_score = 0.0
+    
+    ratios = [len(p) / len(s) if len(s) > 0 else 0 for p, s in zip(predictions, sources)]
+    avg_ratio = sum(ratios) / len(ratios) if ratios else 0
+    
+    return {"SARI": sari_score, "BLEU": bleu.score, "Compression Ratio": avg_ratio}
+
+@torch.no_grad()
+def inference_batch(ae, flow, loader, tokenizer, device, steps=10, save_path="results.txt", use_oneshot=True):
+    ae.eval()
+    flow.eval()
+    stop_id = _pick_stop_id(tokenizer)
+    pad_id = tokenizer.pad_token_id
+    
+    print(f"\n>>> Running Inference on {len(loader.dataset)} examples...")
+    
+    all_sources, all_targets, all_generated = [], [], []
+    scale = getattr(ae, "latent_scale", 10.0) # 兼容逻辑
+    
+    with open(save_path, "w", encoding="utf-8") as f:
+        f.write("Source\tTarget\tGenerated\n")
+        
+        for batch in tqdm(loader, desc="Inferencing"):
+            src_ids = batch['src_ids'].to(device)
+            src_mask = batch['src_mask'].to(device)
+            tgt_ids = batch['tgt_ids'].to(device)
+            B, L = src_ids.shape
+
+            # Encode
+            z_curr = ae.encode(src_ids, src_mask)
+            z_cond = z_curr.clone()
+
+            # Flow Sampling
+            if use_oneshot:
+                t0 = torch.zeros(B, device=device)
+                z_curr = flow(z_curr, t0, condition=z_cond).float()
+            else:
+                dt = 1.0 / steps
+                for i in range(steps):
+                    t_val = i / steps
+                    if t_val >= 0.999: break
+                    t = torch.ones(B, device=device) * t_val
+                    pred_z1 = flow(z_curr, t, condition=z_cond).float()
+                    v = (pred_z1 - z_curr) / (1.0 - t_val + 1e-4) 
+                    z_curr = z_curr + v * dt
+                z_curr = pred_z1 
+            
+            # Decode (Pass 1: Detect Length)
+            full_mask = torch.ones(B, L, device=device)
+            logits1 = ae.decode(z_curr, attention_mask=full_mask)
+            ids1 = logits1.argmax(dim=-1)
+
+            stop_pos = []
+            for i in range(B):
+                pos = _first_pos(ids1[i], stop_id, default=L - 1)
+                stop_pos.append(pos)
+            
+            # Decode (Pass 2: Clean Decode)
+            gen_mask = torch.zeros(B, L, device=device)
+            for i in range(B):
+                gen_mask[i, : stop_pos[i] + 1] = 1.0
+            
+            logits2 = ae.decode(z_curr, attention_mask=gen_mask)
+            ids2 = logits2.argmax(dim=-1)
+            ids2 = ids2.masked_fill(gen_mask == 0, pad_id)
+
+            # Convert to Text
+            src_texts = tokenizer.batch_decode(src_ids, skip_special_tokens=True)
+            tgt_texts = tokenizer.batch_decode(tgt_ids, skip_special_tokens=True)
+            
+            gen_texts = []
+            for i in range(B):
+                end = stop_pos[i] + 1
+                ids_cut = ids2[i, :end]
+                gen_texts.append(tokenizer.decode(ids_cut, skip_special_tokens=True))
+    
+            for s, t, g in zip(src_texts, tgt_texts, gen_texts):
+                s_c = s.replace("\n", " ")
+                t_c = t.replace("\n", " ")
+                g_c = g.replace("\n", " ")
+                f.write(f"{s_c}\t{t_c}\t{g_c}\n")
+                all_sources.append(s_c)
+                all_targets.append(t_c)
+                all_generated.append(g_c)
+                
+    return all_sources, all_targets, all_generated
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ae_ckpt", type=str, default="/mnt/hdfs/user/lixinyu.222/CodeFlow/residual_robust_checkpoints/ae_best.pt", help="Path to pre-trained AE checkpoint")
+    parser.add_argument("--save_dir", type=str, default="residual_robust_checkpoints", help="Directory to save flow checkpoints")
+    parser.add_argument("--use_oneshot", action="store_true", default=True, help="Use one-shot sampling for inference")
+    args = parser.parse_args()
+
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # --- Config ---
+    m_cfg = ModelConfig(
+        encoder_name='../jina-embeddings-v2-base-code',
+        latent_dim=512, 
+        max_seq_len=128
+    )
+    
+    t_cfg = TrainConfig(
+        batch_size=16,
+        num_epochs_flow=35, # 只关注 Flow 的 epoch
+        grad_accum_steps=4,
+        use_amp=False,
+        lr_flow=2e-4
+    )
+    
+    # --- Tokenizer & Data ---
+    tokenizer = AutoTokenizer.from_pretrained(m_cfg.encoder_name,local_files_only=True, trust_remote_code=False)
+    train_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="train")
+    test_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="test")
+
+    # --- Load AE (Pre-trained) ---
+    print(f"\n>>> Loading Pre-trained Autoencoder from {args.ae_ckpt} ...")
+    ae = ReshapedAutoencoder(m_cfg).to(t_cfg.device).float()
+    
+    if not os.path.exists(args.ae_ckpt):
+        raise FileNotFoundError(f"AE checkpoint not found at {args.ae_ckpt}. Please run train_ae.py first.")
+        
+    ae.load_state_dict(torch.load(args.ae_ckpt, map_location=t_cfg.device))
+    
+    # 冻结 AE 的所有参数，Flow 训练时不更新 AE
+    ae.eval()
+    for param in ae.parameters():
+        param.requires_grad = False
+    print(">>> Autoencoder loaded and frozen.")
+
+    if ae.encoder.config.pad_token_id is None:
+        ae.encoder.config.pad_token_id = tokenizer.pad_token_id
+
+    # --- Initialize Flow ---
+    flow = PatchedFlowDiT(m_cfg).to(t_cfg.device).float()
+    
+    # --- Trainer ---
+    trainer = Trainer(
+        ae=ae, 
+        flow=flow, 
+        cfg=t_cfg, 
+        loader=train_loader,
+        pad_id=tokenizer.pad_token_id, 
+        stop_id=_pick_stop_id(tokenizer)
+    )
+
+    # --- Optimizer ---
+    opt_flow = optim.AdamW(flow.parameters(), lr=t_cfg.lr_flow)
+
+    # --- Training Loop ---
+    best_flow_loss = float('inf')
+    print("\n>>> Start Training Flow DiT...")
+    
+    for epoch in range(t_cfg.num_epochs_flow):
+        # 传入 opt_flow 训练 Flow
+        loss = trainer.train_flow(opt_flow)
+        print(f"Flow Epoch {epoch}: Loss {loss:.4f}")
+
+        # Save Best
+        if loss < best_flow_loss:
+            best_flow_loss = loss
+            save_path = os.path.join(args.save_dir, "flow_best.pt")
+            torch.save(flow.state_dict(), save_path)
+            # print(f"  Saved Best Flow to {save_path}")
+            
+        # Save Last
+        torch.save(flow.state_dict(), os.path.join(args.save_dir, "flow_last.pt"))
+    
+    print(f"Flow Training Done. Best Loss: {best_flow_loss:.4f}")
+    
+    # --- Inference / Evaluation ---
+    print("\n>>> Loading Best Flow Checkpoint for Evaluation...")
+    best_flow_path = os.path.join(args.save_dir, "flow_best.pt")
+    if os.path.exists(best_flow_path):
+        flow.load_state_dict(torch.load(best_flow_path, map_location=t_cfg.device))
+    else:
+        print("Warning: Best checkpoint not found, utilizing last epoch weights.")
+
+    print("\n--- Starting Inference ---")
+    sources, targets, gens = inference_batch(
+        ae, flow, test_loader, tokenizer, t_cfg.device, 
+        steps=10, 
+        save_path="wiki_results.tsv",
+        use_oneshot=args.use_oneshot
+    )
+    
+    # Metrics
+    metrics = calculate_metrics(sources, gens, targets)
+    print("\n=== Metrics ===")
+    for k, v in metrics.items():
+        print(f"{k}: {v:.4f}")
+    
+    print(f"\nResults saved to wiki_results.tsv")
+
+if __name__ == "__main__":
+    main()

wikilarge-dataset/.gitattributes

@@ -0,0 +1,55 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.lz4 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Audio files - uncompressed
+*.pcm filter=lfs diff=lfs merge=lfs -text
+*.sam filter=lfs diff=lfs merge=lfs -text
+*.raw filter=lfs diff=lfs merge=lfs -text
+# Audio files - compressed
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+# Image files - uncompressed
+*.bmp filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+# Image files - compressed
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
+wiki.full.aner.ori.train.95.tsv filter=lfs diff=lfs merge=lfs -text

wikilarge-dataset/wiki.full.aner.ori.test.95.tsv

@@ -0,0 +1,192 @@
+Normal	Simple
+his next work saturday follows an especially eventful day in the life of a successful neurosurgeon.	his next work at saturday will be a successful neurosurgeon.
+the tarantula the trickster character spun a black cord and attaching it to the ball crawled away fast to the east pulling on the cord with all his strength.	the tarantuala the trickster spun a black chord and attached it to the ball crawling away fast to the east and pulling the chord with all his strength.
+there he died six weeks later on january.	he died six weeks later on january th.
+they are culturally akin to the coastal peoples of papua new guinea.	their culture is similar to the culture of the coastal peoples of papua new guinea.
+following the drummers are dancers who often play the sogo lrb a tiny drum that makes almost no sound rrb and tend to have more elaborate even acrobatic choreography.	the drummers are dancers and often play the sogo which they tend to have arcobatic choreography.
+the spacecraft consists of two main elements the nasa cassini orbiter named after the italian french astronomer giovanni domenico cassini and the esa huygens probe named after the dutch astronomer mathematician and physicist christiaan huygens.	the spacecraft is having two main parts. one is known as nasa cassini orbiter. it is named after giovanni domenico cassini an italian french astronomer. the other part is known as esa huygens probe. it is named after christiaan huygens. he was a dutch astronomer mathematician and physicist.
+alessandro lrb sandro rrb mazzola lrb born november rrb is an italian former football player.	alessandro mazzola is an italian former football player.
+it was originally thought that the debris thrown up by the collision filled in the smaller craters.	it was thought that the debris thrown up by the collision filled the smaller craters.
+graham attended wheaton college from to when he graduated with a ba in anthropology.	graham attended wheaton college from to and graduated with a ba in anthropology.
+he was also named sportsman of the year by sports illustrated.	sports illustrated named him sportsman of the year in.
+fives is a british sport believed to derive from the same origins as many racquet sports.	fives a british sport came from the same games as many racquet sports.
+for example king bhumibol was born on monday so on his birthday throughout thailand will be decorated with yellow color.	all over thailand the color yellow will be used to celebrate king bhumibal
+nevertheless tagore emulated numerous styles including craftwork from northern new ireland haida carvings from the west coast of canada lrb british columbia rrb and woodcuts by max pechstein.	tagore emulated various styles including craftwork from northern new ireland carvings from the west coast of canada and woodcuts by max pechstein neverthelessly
+she performed for president reagan in s great performances at the white house series which aired on the public broadcasting service.	she did a show for president reagan in s great performances at the white house series which was shown on television on the public broadcasting service.
+perry saturn lrb with terri rrb defeated eddie guerrero lrb with chyna rrb to win the wwf european championship lrb rrb saturn pinned guerrero after a diving elbow drop.	perry saturn and terri defeated eddie guerrero and chyna to win the wwf european championship lrb rrb saturn pinned guerrero after a diving elbow drop.
+she remained in the united states until when she and her husband returned to france.	she stayed in the united states until then she and her husband went to france.
+he also completed two collections of short stories entitled the ribbajack other curious yarns and seven strange and ghostly tales.	he also completed two collections of short storeis. the one title was the ribbajack other curious yarns. the other one was titled as seven strange and ghostly tales.
+at the voyager images ophelia appears as an elongated object the major axis pointing towards uranus.	at the voyager pictures ophelia appears as a stretched object. a stretched object was the major axis. it pointing towards uranus.
+the british decided to eliminate him and take the land by force.	the british decided to put an end to him and take the land by force.
+some towns on the eyre highway in the south east corner of western australia between the south australian border almost as far as caiguna do not follow official western australian time.	there are some towns in western australia that do not follow official western australian time.
+in architectural decoration small pieces of colored and iridescent shell have been used to create mosaics and inlays which have been used to decorate walls furniture and boxes.	small pieces of colored and shiny shell has been used to decorate walls furniture and boxes.
+the other incorporated cities on the palos verdes peninsula include rancho palos verdes rolling hills estates and rolling hills.	rancho palos verdes rolling hills estates and rolling hills are three cities on the palos verdes peninsula.
+fearing that drek will destroy the galaxy clank asks ratchet to help him find the famous superhero captain qwark in an effort to stop drek.	clank asks ratchet to help him find the famous superhero captain qwark because he is afraid that drek will try to destroy the galaxy and wants to stop him.
+he advocates applying a user centered design process in product development cycles and also works towards popularizing interaction design as a mainstream discipline.	he favors product development cycles that features an easy to use design process and works towards bringing interaction design into mainstream popuarity.
+it is theoretically possible that the other editors who may have reported you and the administrator who blocked you are part of a conspiracy against someone half a world away they ve never met in person.	it is possible that the other editors who may have reported you is a part of the conspiracy. similarly the administrator who blocked you may also be a part of the conspiracy. the conspiracy is against someone they have not met in prison.
+working group i assesses scientific aspects of the climate system and climate change.	working group i makes note of climate system and climate change
+formal minor planet designations are number name combinations overseen by the minor planet center a branch of the iau.	formal minor planet designations are number name combinations overseen by the minor planet center.
+as a result although many mosques will not enforce violations both men and women when attending a mosque must adhere to these guidelines.	although many mosques will not enforce rules both men and women when there must follow these rules
+mariel of redwall is a fantasy novel by brian jacques published in.	mariel of redwall is a fiction novel in the category of fantasy by author brian jacques published in.
+ryan prosser lrb born july rrb is a professional rugby union player for bristol rugby in the guinness premiership.	ryan prosser who was born on july is a professional rugby union player he has played for briston rugby in the guinness premiership.
+like previous assessment reports it consists of four reports three of them from its working groups.	the assessment report contiains four reports just like previous reports and three of them are from working groups.
+this stamp remained the standard letter stamp for the remainder of victoria s reign and vast quantities were printed.	this stamp stayed the standard letter stamp for the rest of victoria s reign and many were printed.
+the international fight league was an american mixed martial arts lrb mma rrb promotion billed as the world s first mma league.	the world s first mma league was the international fight league and american mixed martial arts.
+giardia lamblia lrb synonymous with lamblia intestinalis and giardia duodenalis rrb is a flagellated protozoan parasite that colonises and reproduces in the small intestine causing giardiasis.	giardia lamblia is a flagellated protozoan parasite that colonises and reproduces in the small intestine causing giardiasis.
+this was the area east of the mouth of the vistula river later sometimes called prussia proper.	this area which later was sometimes called prussia proper was east of the place where the vistula river begins.
+after graduation he returned to yerevan to teach at the local conservatory and later he was appointed artistic director of the armenian philarmonic orchestra.	fter graduation he came back to yerevan to teach at the local conservatory and then he was appointed artistic director of the armenian philarmonic orchestra.
+weelkes was later to find himself in trouble with the chichester cathedral authorities for his heavy drinking and immoderate behaviour.	weelkes later found himself in trouble with the chichester cathedral authorities for his heavy drinking and immoderate behavior.
+he is also a member of another jungiery boyband club.	he is also a member of club which is another jungiery boyband.
+the apostolic tradition attributed to the theologian hippolytus attests the singing of hallel psalms with alleluia as the refrain in early christian agape feasts.	the apostolic tradition connected to the scientist hippolytus who is an expert in theology starts the singing of hallel psalms with alleluia as the repeated line in early christian lovable and wonderful festivals.
+it was the first asteroid to be discovered by a spacecraft.	it was the first asteroid discovered by a spacecraft.
+it continues as the bohemian switzerland in the czech republic.	it is still called as the bohemian switzerland in the czech republic.
+this leads to consumer confusion when lrb rrb bytes is referenced as mb lrb megabyte rrb instead of mib.	the consumer gets confused when bytes is called mb instead of mib.
+the incident has been the subject of numerous reports as to ethics in scholarship.	the incident has been the subject of numerous reports regarding scholarship ethics.
+they are castrated so that the animal may be more docile or may put on weight more quickly.	they are castrated so that the animal is docile or may put on weight quickly.
+seventh sons have strong knacks lrb specific magical abilities rrb and seventh sons of seventh sons are both extraordinarily rare and powerful.	seventh sons have strong knacks lrb specific magical abilities rrb and he is extraordinarily rare and powerful.
+benchmarking conducted by passmark software highlights the version s second install time second scan time and mb memory utilization.	passmark software tested standards of version and the highlights are second install time second scan time and mb memory utilization.
+if there are no strong land use controls buildings are built along a bypass converting it into an ordinary town road and the bypass may eventually become as congested as the local streets it was intended to avoid.	if there are no strong land use controls the bypass as a result may become congested. this is because buildings are built along a bypass converting it into an ordinary town road. the byepass is intended to avoid such congestion.
+it is also a starting point for people wanting to explore cooktown cape york peninsula and the atherton tableland.	it is a starting point for people wanting to explore cooktown cape york peninsula and atherton tableland.
+bruises often induce pain but are not normally dangerous.	bruises often hurt but are not normally dangerous.
+tickets can be retailed for national rail services the docklands light railway and on oyster card.	tickets can be retailed for national rail services and the docklands light railway on oyster card.
+the historical method comprises the techniques and guidelines by which historians use primary sources and other evidence to research and then to write history.	in writing history there is a method called the historical method which uses primary sources and other evidence to research the historical events.
+the sheer weight of the continental icecap sitting on top of lake vostok is believed to contribute to the high oxygen concentration.	the lake vostok has a very large weight of the continental icecap on its waters. the high oxygen concentration of the lake water may because of this icecap.
+aliteracy lrb sometimes spelled alliteracy rrb is the state of being able to read but being uninterested in doing so.	aliteracy is being able to read but uninterested in read.
+mifepristone is a synthetic steroid compound used as a pharmaceutical.	pharmaceutical has used a mifepristone is a synthetic steroid
+shortly after attaining category status the outer convection of the hurricane became ragged.	shortly after reaching category status the outer convection of the hurricane became worn out
+the equilibrium price for a certain type of labor is the wage.	the balanced price for any kind of labor is called a wage.
+convinced that the grounds were haunted they decided to publish their findings in a book an adventure lrb rrb under the pseudonyms of elizabeth morison and frances lamont.	the authors using pseudonyms elizabeth morison and frances lamont published the book an adventure in written about particular hauntings.
+he left a detachment of troops to garrison the newly conquered region.	he left troops to garrison the newly conquered region.
+the depression moved inland on the th as a circulation devoid of convection and dissipated the next day over brazil where it caused heavy rains and flooding.	on th depression moved inland as convection less circulation and after weakening over brazil it caused heavy rains and flooding.
+the characters are foul mouthed extensions of their earlier characters pete and dud.	the characters speak bad language of their earlier characters pete and dud
+in culver ran for iowa secretary of state and was victorious.	in culver successfully ran for iowa secretary of state.
+in mark messier took the hart over ray bourque by a margin of two votes the difference being a single first place vote.	in mark messier took hart over ray bourque by two votes the difference being a first place vote.
+shade sets the main plot of the novel in motion when he impetuously defies that law and inadvertently initiates a chain of events that leads to the destruction of his colony s home forcing their premature migration and his separation from them.	the main plot of the novel is when shade defies the law and sets off a chain of events that lead to the destuction of his colony s home and becoming separated from them.
+the amazon basin is the part of south america drained by the amazon river and its tributaries.	the amazon basin is the part of south america drained by the amazon river and those who pay tribute to it.
+the two former presidents were later separately charged with mutiny and treason for their roles in the coup and the gwangju massacre.	the two former presidents were later charged each on their own with mutiny and treason for their roles in the coup and the gwangju massacre.
+moderate to severe damage extended up the atlantic coastline and as far inland as west virginia.	there was moderate to severe damage all the way up the atlantic coastline and as far inland as west virginia.
+because the owner tends to be unaware these computers are metaphorically compared to zombies.	these computers are metaphorically compared to zombies as the owner was not concious about it
+for example the stylebook of the associated press is updated annually.	the stylebook of the associated press is updated yearly.
+the four canonical texts are the gospel of matthew gospel of mark gospel of luke and gospel of john probably written between ad and lrb see also the gospel according to the hebrews rrb.	gospels matthew mark luke and john were most likey written after christ
+development stable releases are rare but there are often subversion snapshots which are stable enough to use.	development stable releases are rare but there are often subversion snapshots lrb an informal photograph taken quickly rrb which are stable enough to use.
+cogeneration lrb also combined heat and power chp rrb is the use of a heat engine or a power station to simultaneously generate both electricity and useful heat.	cogeneration is a combination of heart and power to simultaneously generate both electricity and useful heat.
+on occasion the male den master will also allow a second male into the den the reason for this is unclear.	on opportunity the male den master will let a second male inside of the den the basis for this is poorly explained.
+below are some useful links to facilitate your involvement.	below are some useful links to help your involvement.
+he served as the prime minister of egypt between and and again from and.	he served as prime minster of egypt from and as well as from through.
+she was left behind lrb explanations for this vary rrb when the rest of the nicole os were moved to the mainland.	people have different thinking about why she was left behind when the nicolenos were moved to the main part of the country.
+chauvin was embarrassed to receive his award and initially indicated that he may not accept it.	chauvin was embarrassed to get his award and at first said that he may not accept it.
+later esperanto speakers began to see the language and the culture that had grown up around it as ends in themselves even if esperanto is never adopted by the united nations or other international organizations.	later esperanto speakers started to see the language and culture that had grown up around it as ends in themselves though esperanto is never accepted by the united nations of other international organizations.
+dry air wrapping around the southern periphery of the cyclone eroded most of the deep convection by early on september.	early september dry air wrapping around the southern area of the cyclone caused most of the heat to leave.
+a few animals have chromatic response changing color in changing environments either seasonally lrb ermine snowshoe hare rrb or far more rapidly with chromatophores in their integument lrb the cephalopod family rrb.	some animals change color when their environments change a process called chromatic response either seasonally as with ermine and snowshoe hare or far more rapidly with chromoa tophonres in theri integument lrb the cephalapod family. rrb 
+this closely resembles the unix philosophy of having multiple programs each doing one thing well and working together over universal interfaces.	this looks like the unix idea of having several programs with each doing one thing and working together.
+he came from a musical family his mother larue was an administrative assistant and singer and his father keith brion was a band director at yale.	his was a musical family as his mother larue was a secretary and singer while his father keith brion was a band director at yale.
+the largest populations of mennonites are in canada democratic republic of congo and the united states but mennonites can also be found in tight knit communities in at least countries on six continents or scattered amongst the populace of those countries.	the largest populations of mennonites are in canada democratic republic of congo and the united states. mennonites also live in close communities in at least countries on six continents or scattered throughout the populations of those countries.
+acanthopholis s armour consisted of oval plates set almost horizontally into the skin with spikes protruding from the neck and shoulder area along the spine.	acanthopholis s armour was made up of oval plates that were put into the skin lengthwise and had spikes that jutted out from the neck and shoulder area across the spine.
+conversely bills proposed by the law commission and consolidation bills start in the house of lords.	the bills proposed by the law commission and consolidation bills start in the house of lords contrarily.
+reflection nebulae are usually blue because the scattering is more efficient for blue light than red lrb this is the same scattering process that gives us blue skies and red sunsets rrb.	reflection nebulae are commonly blue because the scattering is more powerful for blue light than red lrb this is the same reason for the sky appears in blue and the sunset in red colors
+macgruber starts asking for simple objects to make something to defuse the bomb but he is later distracted by something lrb usually involving his personal life rrb that makes him run out of time.	macgruber asked for many items to help shut the bomb off but he was distracted and ran out of time.
+the pad called for the resignation of the governments of thaksin shinawatra samak sundaravej and somchai wongsawat whom the pad accused of being proxies for thaksin.	the pad called for thakin shinaatra samak sundaravej and somchai to step down as government leaders because the pad considered them to be used by thaksin.
+while at kahn he was chief architect for the fisher building in.	he was a chief architect fot fisher building in when he was at kahn
+he excuses himself because he has to leave for rehearsal and he and dr. sch n leave.	he excuses himself because he has to leave for rehearsal dr.dr. sch n leave
+the sheppard line currently has fewer users than the other two subway lines and shorter trains are run.	the sheppard line not only has fewer users than the other two subway lines it also runs shorter trains.
+it has a capacity of making it the largest stadium in europe and the eleventh largest in the world.	it can seat. which makes it the largest stadium in europe and the eleventh largest in the world.
+in december ten boom was honored as one of the righteous among the nations by the state of israel.	in december ten boon was honored as part of the righteous amoung the nations by the state of israel.
+terms such as undies for underwear and movie for moving picture are oft heard terms in english.	words like undies movie are oft heard terms in english.
+jurisdiction draws its substance from public international law conflict of laws constitutional law and the powers of the executive and legislative branches of government to allocate resources to best serve the needs of its native society.	power moves towards its material from community national boundaries opposition of laws organizational law and ability of the administrative and creative offshoots of government to assign support to best serve the needs of its native society.
+despite this farrenc was paid less than her male counterparts for nearly a decade.	even so farrenc was paid less than her male peers for nearly years.
+gumbasia was created in a style vorkapich taught called kinesthetic film principles.	vorkapich taught gumbasia in a style called kinesthetic film principles.
+the lawyer brandon lrb waise lee rrb became his idol and mk sun grew up to be a lawyer.	the lawyer brandon lrb waise lee rrb was his idol as mk sun grew up to be a lawyer.
+military career donaldson enlisted in the australian army on june.	donaldson inlisted in australia s army on june to start his military career.
+the kindle features level grayscale display improved battery life percent faster page refreshing a text to speech option to read the text aloud and overall thickness reduced from. to. inches lrb. millimeters rrb.	the kindle features grayscale display improved battery life and overall thickness reduced.
+yoghurt or yogurt is a dairy product produced by bacterial fermentation of milk.	yoghurt or yogurt is a milk based food made by bacterial fermentation of milk.
+seventy five defencemen are in the hall of fame more than any other current position while only goaltenders have been inducted.	out of seventy five defencemen in the hall of fame only goaltenders have been inducted.
+alternative views on the subject have been proposed throughout the centuries lrb see below rrb but all were rejected by mainstream christian bodies.	different views on the subject have been brought up over the centuries lrb see below rrb but all were rejected by mainstream christian bodies 
+the album however was banned from many record stores nationwide.	the album is banned from many record stores nationwide.
+the company opened twice as many canadian outlets as mcdonald s wendy s confirms tim hortons ipo by march ottawa business journal december and system wide sales also surpassed those of mcdonald s canadian operations as of.	the company opened two times as many restaurants in canada as mcdonald s wendy s confirms tim hortons ipo by march ottawa business journal december and sales throughout the company were also greater than those of mcdonald s canadian business as of.
+he won the presidential election held on march with. of the popular vote.	he conquered the presidential poll on march with. of the popular vote.
+in she was the only female entertainer allowed to perform in saudi arabia.	as a female entertainer she alone was allowed to perform in saudi arabia during 
+offenbach s numerous operettas such as orpheus in the underworld and la belle h l ne were extremely popular in both france and the english speaking world during the s and s.	offenbach s a great number of operettas such as orpheus in the underworld and la beautiful woman helene were greatly pleasing to all in both france and the english talking earth during the s and s
+roof tiles dating back to the tang dynasty with this symbol have been found west of the ancient city of chang an lrb modern day xian rrb.	roof tiles during tang dynasty with this symbol have been found west of the ancient city of chang an or modern day xian.
+by most accounts the instrument was nearly impossible to control.	it was nearly impossible to control the instrument by most accounts.
+characteristics radar observations indicate a fairly pure iron nickel composition.	radar testing shows composition of mostly iron nickel.
+lo che harbours the installations of onyx the swiss interception system for electronic intelligence gathering.	lo che harbours of onyx is the swiss interception system for electronic intelligence gathering.
+a matchbook is a small cardboard folder lrb matchcover rrb enclosing a quantity of matches and having a coarse striking surface on the exterior.	a matchbook is a small cardboard folder lrb or matchcover rrb that holds some matches and has a rough area on the outside.
+she was among the first doctors to object to cigarette smoking around children and drug use in pregnant women.	she was one of the first doctors that said cigarette smoking near children and drug use in pregnant women was not safe.
+defiantly she vowed to never renounce the commune and dared the judges to sentence her to death.	she refused to give up the commune and prefered the death sentence
+oel manga series graystripe s trilogy there is a three volume original english language manga series following graystripe between the time that he was taken by twolegs in dawn until he returned to thunderclan in the sight.	oel manga occuring in sequence graystripe s trilogy. there s a three quantity volume earliest english language manga series following graystripe linking the time that it was accepted by twolegs in dawn up until he came back to thunderclan in the sight.
+samovar porter lrb rrb p. syrians did not congregate in urban enclaves many of the immigrants who had worked as peddlers were able to interact with americans on a daily basis.	samovar porter lrb rrb p. syrians did not get together in city groups many of the immigrants who had worked as sellers on the street were able to talk with americans every day.
+he was also famous for his prints book covers posters and garden metalwork furniture.	he is famous for prints book covers posters and garden metalwork furniture.
+during childhood she suffered from collapsed lungs twice she had pneumonia times a year a ruptured appendix and had a tonsillar cyst.	for two times she had lung disorder when she was a child. she was also suffered from pneumonia to times a year. she was also affected by appendix disorder and had a tonsillar cyst. all these happened during her childhood period.
+small value inductors can also be built on integrated circuits using the same processes that are used to make transistors.	both small value inductors and transistors can be built on integrated circuits.
+no skater has yet accomplished a quadruple axel in competition.	quadruple axel at a competition is yet to be fulfilled by any skater
+from the telephone exchange the port jackson district commandant could communicate with all military installations on the harbour.	by use of the telephone exchange the port jackson district commandant could talk to all military installations on the harbour.
+however even to those who enter the prayer hall of a mosque without the intention of praying there are still rules that apply.	however even to those who enter the prayer hall of a mosque without the purpose of praying there are some rules applicable to them.
+it is described as pointed in the face and about the size of a rabbit.	it is about the size of a rabbit and has a pointed face.
+human skin hues can range from very dark brown to very pale pink.	the colors of human skin can be very dark brown or very pale pink or anywhere in between.
+nupedia was founded on march under the ownership of bomis inc a web portal company.	bomis inc a web portal company founded nupedia on march.
+notable features of the design include key dependent s boxes and a highly complex key schedule.	notable features of the design include s boxes which is a highly complex key schedule.
+the primavera is a painting by the italian renaissance painter sandro botticelli c..	painted around the primavera is a painting by the italian renaissance painter sandro botticelli.
+new south wales s largest city and capital is sydney.	largest city new south wales and its capital is sydney.
+the polymer is most often epoxy but other polymers such as polyester vinyl ester or nylon are also sometimes used.	polymers such as polyester vinylester or nylon are also sometimes used as epoxies.
+stands were eventually added behind each set of goals during the s and s as the ground began to be modernised.	during the s and s the ground got more modern and stands were eventually added behind each set of goals.
+a bastion on the eastern approaches was built later.	later a bastion was built on the eastern approaches.
+events europe july battle of stiklestad lrb norway rrb olav haraldsson loses to his pagan vassals and is killed in the battle.	among events that happened in europe on july. was the battle of stiklesstand in norway in which olav haraldsson lost his pagan subjects and his life.
+others have theorized that tresca was eliminated by the nkvd as retribution for criticism of the stalin regime of the soviet union.	others have made theories that tresca was took away by the nkvd as retribution for views put forward as to errors of the stalin system of things of the soviet union.
+schuschnigg immediately responded publicly that reports of riots were false.	schuschnigg responded publicly that reports of riots were false.
+depending on the context another closely related meaning of constituent is that of a citizen residing in the area governed represented or otherwise served by a politician sometimes this is restricted to citizens who elected the politician.	based on the given situation one more nearest meaning of constituent would be a citizen residing in the area governed represented or otherwise served by a politician at times this is limited to the citizens who elected the politician.
+wario land the wario land series is a platforming series that started with wario land super mario land a spin off of the super mario land series.	the platform series wario land from the wario land series stared with the super mario land series.
+these attacks may have been psychological in origin rather than physical.	these attacks may have been psychological in origin not physical.
+furthermore spectroscopic studies have shown evidence of hydrated minerals and silicates which indicate rather a stony surface composition.	furthermore spectroscopic studies have shown proof of hydrated minerals and silicates which points to rather a stony surface material.
+she became the authoritative editor of her husband s works for breitkopf und h rtel.	she became the editor of her husband s works for breitkopf und hartel.
+mercury is similar in appearance to the moon it is heavily cratered with regions of smooth plains has no natural satellites and no substantial atmosphere.	mercury is similar in appearance to the moon with heavily cratered regions of smooth plains and no natural satellites or substantial atmosphere.
+geography the town lies in the limmat valley between baden and z rich.	geographically the town lies in the limmat valley between baden and zurich.
+these ideally provide excellent habitat for chinkara hog deer and blue bull.	ideally these make an excellent breeding ground for chinkara hog deer and blue bull.
+after the sena dynasty dhaka was successively ruled by the turkish and afghan governors descending from the delhi sultanate before the arrival of the mughals in.	before the mughals arrived in after the sena dynasty dhaka was ruled for a long time by turkish and afghan governors that descended from the delhi sultanate.
+for rowling this scene is important because it shows harry s bravery and by retrieving cedric s corpse he demonstrates selflessness and compassion.	for rowling this incident is important becaus it shows harry s courage and by regaining cedric d corpse he establishes no concern for oneself and sympathy.
+on june he and fellow raf members jan carl raspe and holger meins were apprehended after a lengthy shootout in frankfurt.	he and fellow raf members jan carl raspe and holger meins were taken hold after a lengthy shootout. it was in frankfurt on june.
+together they formed new music manchester a group committed to contemporary music.	they formed the new music manchester band and sang contemporary
+the compact and intense hurricane caused extreme damage in the upper florida keys as a storm surge of approximately to feet affected the region.	small but intense the hurricane caused a lot of damage in the upper florida keys when a surge of nearly to feet hit the area.
+it is now the site of meher baba s samadhi lrb tomb shrine rrb as well as facilities and accommodations for pilgrims.	it is now the meher baba s tomb shrine and a place for pilgrims.
+salem is a city in essex county massachusetts united states.	salem is a city in essex county massachusetts.
+forty nine species of pipefish and nine species of seahorse have been recorded.	forty nine species of pipefish and nine species of seahorse are recorded.
+therefore these pdfs can not be distributed without further manipulation if they contain images.	if any of these pdfs contain pictures then they require additional processing before they can be issued
+heavy rain fell across portions of britain on october causing localized accumulation of flood waters.	heavy rain fell across britain on october causing accumulation of flood waters.
+ohio state s library system encompasses twenty one libraries located on its columbus campus.	ohio state s library system has twenty one libraries located on its campus.
+in other developments both iceland and greenland accepted the overlordship of norway but scotland was able to repulse a norse invasion and broker a favorable peace settlement.	both iceland and greenland accepted the ruler of norway but scotland was able to prevent a norwegian invasion and a negotiate a peace settlement.
+the singles from the album included by the way the zephyr song ca n t stop dosed and universally speaking.	some singles from the album are by the way the zephyr song ca n t stop dosed and universally speaking.
+in april minix became free open source software under a permissive free software licence but by this time other operating systems had surpassed its capabilities and it remained primarily an operating system for students and hobbyists.	in april minix became free or open source software under a non restrictive free software licence but by this time other operating systems had exceeded its capabilities and it continued to be mainly an operating system for students and hobbyists.
+the body color varies from medium brown to gold ish to beige white and occasionally is marked with dark brown spots especially on the limbs.	the body color varies from medium brown to goldish to beige white and sometimes is marked with dark brown spots.
+the latter provided audiences with the sort of information later provided by intertitles and can help historians imagine what the film may have been like.	the latter which gave audiences the same sort of information later audience members would gett from subtitles can help historians imagine what the film may have been like.
+that is because real estate businesses and other assets in the underground economies of the third world can not be used as collateral to raise capital to finance industrial and commercial expansion.	that is beacuase real estate businesses and other assets in the underground economies of the third world can not be used as collateral to raise capital.
+ned and dan advanced to the police camp ordering them to surrender.	ned and dan told by the police to surrender.
+a mutant is a type of fictional character that appears in comic books published by marvel comics.	mutants are fictional characters from the x men comic books published by marvel.
+the sat reasoning test lrb formerly scholastic aptitude test and scholastic assessment test rrb is a standardized test for college admissions in the united states.	the sat is a standardized test for college admissions in the united states.
+some reports read that various factors increase the likelihood of both paralysis and hallucinations.	some reports said that various things make it more possible to have paralysis and hallucinations.
+his sentence was transportation to australia for seven years.	his sentence was carried to australia for seven years.
+her notorious friendship with the russian mystic grigori rasputin was also an important factor in her life.	her well known relation with the russian mystic grigori rasputin was additionaly an important number in her life.
+the term dorsal refers to anatomical structures that are either situated toward or grow off that side of an animal.	the word dorsal means any body part that grows off that side of an animal or that grows toward that side of an animal.
+the term protein itself was coined by berzelius after mulder observed that all proteins seemed to have the same empirical formula and might be composed of a single type of lrb very large rrb molecule.	the term protein was made by berzelius after mulder
+after the jerilderie raid the gang laid low for months evading capture.	after the jerilderie raid the gang laid low for months.
+in an extension was added curving north from union station below university avenue and queen s park to near bloor street where it turned west to terminate at st. george and bloor streets.	an extension was added in which curved north from union station below university avenue and queens park reaching nearly to bloor street and ending on the west side at st. george and bloor streets.
+it is located on an old portage trail which led west through the mountains to unalakleet.	it s near an old portage trail that led west to unalakleet through the mountains.
+people with cardiomyopathy are often at risk of arrhythmia or sudden cardiac death or both.	arrhythmia or heart beat disorder and sudden cardiac arrest are often associated with cardiomyopathy. cardiomyopathy is deterioration of heart muscle and persons with this disease may subject to arrhythmia or sudden cardiac arrest. sometimes both may happen at once.
+as the largest sub region in mesoamerica it encompassed a vast and varied landscape from the mountainous regions of the sierra madre to the semi arid plains of northern yucat n.	as the largest sub region in mesoamerica it is a vast and varied landscape from the mountainous regions of the sierra madre to the plains of yucatan.
+google subsequently made the comic available on google books and their site and mentioned it on its official blog along with an explanation for the early release.	google made the comic available on google books mentioned it on their blog explaining the early release.
+the book political economy was published in but had limited classroom adoption.	the book political economy was published in but was not used in many classrooms.
+he toured with the ipo in the spring of for their first ever performance in the soviet union with concerts in moscow and leningrad and toured with the ipo again in performing in china and india.	for their first ever performance in the soviet union he toured with the ipo spring from to with china and india
+napoleonic wars austrian general mack surrenders his army to the grand army of napoleon at ulm reaping napoleon over prisoners and inflicting casualties on the losers.	austrian general mack surrenders his army to grand army of napoleon at ulm.
+it has long been the economic centre of northern nigeria and a centre for the production and export of groundnuts.	it has long been the economic centre of norther nigeria along with the center for production and export of groundnuts.
+a majority of south indians speak one of the five dravidian languages kannada malayalam tamil telugu and tulu.	most south indians speak one of the five dravidian languages kannada malayalam tamil telugu and tulu.
+meteora earned the band multiple awards and honors.	 meteora won many awards and honors for the band.
+in the th century slavs started to move into the area.	in the th century slaves started to move in the area.
+winchester is a city in scott county illinois united states.	winchester is a city located in scott county illinois united states
+out of participants in the national casting she was chosen among the candidates to appear on the tv show.	out of participants in the national casting and it was selected the candidates to appear on the tv show.
+the latter device can then be designed and used in less stringent environments.	the device can be designed for use in less exact environments.
+gimnasia hired first famed colombian trainer francisco maturana and then julio c sar falcioni but both had limited success.	gimnasia hired francisco maturana a columbian trainer and then julio c sar falcioni but they were not very successful.
+brighton is a city in washington county iowa united states.	brighton is a city in washington county iowa.
+pauline returned in the game boy remake of donkey kong in and later mario vs. donkey kong march of the minis in although the character is now described as mario s friend.	pauline turned in the game boy remake of donkey kong in and later mario vs. donkey kong march of the minis in even though the character is now mario s friend.
+his real date of birth was never recorded but it is believed to be a date between and.	since his actual date of birth was not recorded it is believed to be between.
+this quantitative measure indicates how much of a particular drug or other substance lrb inhibitor rrb is needed to inhibit a given biological process lrb or component of a process i.e. an enzyme cell cell receptor or microorganism rrb by half.	this quantitative measure indicates how much of a drug or other substance is needed to inhibit a biological process by half.
+there he had one daughter later baptized as mary ann fisher power to ann lrb e rrb power.	he had one daughter named mary ann fisher power who was later baptized to ann lrb e rrb power.
+during an interview edward gorey mentioned that bawden was one of his favorite artists lamenting the fact that not many people remembered or knew about this fine artist.	during an interview edward gorey said that bawden was one of his favorite artists and is saddened by the fact that not many people remembered or knew about this fine artist.
+gable also earned an academy award nomination when he portrayed fletcher christian in s mutiny on the bounty.	gable also earned an academy award nomination for his portrayal of fletcher christian in the film mutiny on the bounty.

wikilarge-dataset/wiki.full.aner.ori.train.95.tsv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3173629a89d5b4a7311262bc7c7fb2ec480c400330410e32730a75afe365dea0
+size 36251904