Upload 9 files

detree/cli/database.py

@@ -0,0 +1,117 @@
+"""Generate clustered prototype databases from embeddings."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Iterable, Optional
+
+import faiss
+import numpy as np
+import torch
+
+
+class GPUKMeansClusterer:
+    def __init__(self, dim: int, n_clusters: int = 500, n_iter: int = 20, n_gpu: int = 1):
+        self.clus = faiss.Clustering(dim, n_clusters)
+        self.clus.verbose = True
+        self.clus.niter = n_iter
+        self.dim = dim
+        self.n_clusters = n_clusters
+        self.clus.update_index = True
+
+        res = [faiss.StandardGpuResources() for _ in range(n_gpu)]
+        flat_config = []
+        for i in range(n_gpu):
+            cfg = faiss.GpuIndexFlatConfig()
+            cfg.useFloat16 = False
+            cfg.device = i
+            flat_config.append(cfg)
+
+        if n_gpu == 1:
+            self.index = faiss.GpuIndexFlatL2(res[0], self.dim, flat_config[0])
+        else:
+            indexes = [faiss.GpuIndexFlatL2(res[i], self.dim, flat_config[i]) for i in range(n_gpu)]
+            self.index = faiss.IndexReplicas()
+            for sub_index in indexes:
+                self.index.addIndex(sub_index)
+
+    def fit(self, embeddings_np: np.ndarray) -> np.ndarray:
+        self.index.reset()
+        self.clus.train(embeddings_np, self.index)
+        centroids = faiss.vector_float_to_array(self.clus.centroids)
+        centroids = centroids.reshape(self.n_clusters, self.dim)
+        return centroids
+
+
+def gen_data(dict_data):
+    embeddings = dict_data["embeddings"]
+    labels = dict_data["labels"]
+    ids = dict_data["ids"]
+    classes = dict_data["classes"]
+    return embeddings, labels, ids, classes
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Cluster embeddings into prototype databases using GPU K-Means.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--database", type=Path, required=True, help="Input embedding database (.pt).")
+    parser.add_argument("--output", type=Path, required=True, help="Output path for the clustered database.")
+    parser.add_argument("--clusters", type=int, default=10000)
+    parser.add_argument("--dimension", type=int, default=1024)
+    parser.add_argument("--iterations", type=int, default=100)
+    parser.add_argument("--gpus", type=int, default=1)
+    parser.add_argument("--human-class-name", type=str, default="human", help="Label representing humans in the class list.")
+    return parser
+
+
+def cluster_database(args: argparse.Namespace) -> None:
+    data_emb, data_labels, data_ids, data_classes = gen_data(torch.load(args.database))
+    human_idx = data_classes.index(args.human_class_name)
+    datapos = (data_labels == human_idx).long()
+    pos2cnt = {0: args.clusters, 1: args.clusters}
+    pos2name = {0: ["llm"], 1: ["human"]}
+
+    datapos_np = datapos.cpu().numpy()
+    kmeans = GPUKMeansClusterer(args.dimension, n_clusters=args.clusters, n_iter=args.iterations, n_gpu=args.gpus)
+    all_centers = {}
+    save_labels = None
+    for key in data_emb:
+        now_emb = data_emb[key].float().cpu().numpy()
+        all_center = []
+        all_labels = []
+        for pos in pos2cnt:
+            pos_emb = now_emb[datapos_np == pos]
+            pos_center = kmeans.fit(pos_emb)
+            all_center.append(pos_center)
+            all_labels.append(np.full((pos_center.shape[0],), pos))
+        all_center = np.concatenate(all_center, axis=0)
+        all_labels = np.concatenate(all_labels, axis=0)
+        all_center = torch.from_numpy(all_center).to(dtype=torch.bfloat16)
+        all_labels = torch.from_numpy(all_labels).to(dtype=torch.long)
+        all_centers[key] = all_center
+        save_labels = all_labels
+
+    save_ids = torch.arange(save_labels.shape[0], dtype=torch.long)
+    classes = [None] * len(pos2name.keys())
+    for pos in pos2name:
+        classes[pos] = ','.join(pos2name[pos])
+
+    emb_dict = {"embeddings": all_centers, "labels": save_labels, "ids": save_ids, "classes": classes}
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(emb_dict, args.output)
+    print(f"All centers saved to: {args.output}")
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    args = parser.parse_args(argv)
+    cluster_database(args)
+
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["build_argument_parser", "cluster_database", "main"]

detree/cli/embeddings.py

@@ -0,0 +1,200 @@
+"""Embedding generation CLI for DETree."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Iterable, Literal, Optional
+
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from lightning import Fabric
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from detree.model.text_embedding import TextEmbeddingModel
+from detree.utils.dataset import SCLDataset, load_datapath
+
+
+def infer(passages_dataloader, fabric, tokenizer, model, args):
+    if fabric.global_rank == 0:
+        passages_dataloader = tqdm(passages_dataloader)
+        all_ids, all_embeddings, all_labels = [], {}, []
+        for layer in args.need_layer:
+            all_embeddings[layer] = []
+    with torch.no_grad():
+        for batch in passages_dataloader:
+            text, label, write_model, ids = batch
+            encoded_batch = tokenizer.batch_encode_plus(
+                text,
+                return_tensors="pt",
+                max_length=args.max_length,
+                padding="max_length",
+                truncation=True,
+            )
+            encoded_batch = {k: v.cuda() for k, v in encoded_batch.items()}
+            embeddings = model(encoded_batch, hidden_states=True)
+            embeddings = fabric.all_gather(embeddings).view(-1, embeddings.size(-2), embeddings.size(-1))
+            label = fabric.all_gather(write_model).view(-1)
+            ids = fabric.all_gather(ids).view(-1)
+            if fabric.global_rank == 0:
+                embeddings = F.normalize(embeddings, dim=-1).cpu().to(torch.bfloat16)
+                for layer in args.need_layer:
+                    all_embeddings[layer].append(embeddings[:, layer, :].clone())
+                all_ids.extend(ids.cpu().tolist())
+                all_labels.extend(label.cpu().tolist())
+            del embeddings, label, ids
+    if fabric.global_rank == 0:
+        for layer in args.need_layer:
+            all_embeddings[layer] = torch.cat(all_embeddings[layer], dim=0)
+        return torch.tensor(all_ids), all_embeddings, torch.tensor(all_labels)
+    return [], [], []
+
+
+def stable_long_hash(input_string: str) -> int:
+    import hashlib
+
+    hash_object = hashlib.sha256(input_string.encode())
+    hex_digest = hash_object.hexdigest()
+    int_hash = int(hex_digest, 16)
+    return int_hash & ((1 << 63) - 1)
+
+
+def load_data(split: Literal["train", "test", "extra"], include_adversarial: bool, fp: Path) -> pd.DataFrame:
+    if split not in ("train", "test", "extra"):
+        raise ValueError("`split` must be one of (\"train\", \"test\", \"extra\")")
+
+    fname = f"{split}.csv" if include_adversarial else f"{split}_none.csv"
+    fp = fp / fname
+    return pd.read_csv(fp)
+
+
+class PassagesDataset(Dataset):
+    def __init__(self, data):
+        self.passages = []
+        for item in data:
+            if item["attack"] not in ("none", "paraphrase") and stable_long_hash(item["generation"]) % 10 < 5:
+                continue
+            self.passages.append(item)
+        classes = sorted({item["model"] for item in data})
+        self.classes = list(classes)
+        self.human_id = self.classes.index("human")
+
+    def __len__(self):
+        return len(self.passages)
+
+    def __getitem__(self, idx):
+        data_now = self.passages[idx]
+        text = data_now["generation"]
+        model = self.classes.index(data_now["model"])
+        label = int(model == self.human_id)
+        ids = stable_long_hash(text)
+        return text, int(label), int(model), int(ids)
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Generate embedding databases for DETree evaluators",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--device-num", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--num-workers", type=int, default=8)
+    parser.add_argument("--max-length", type=int, default=512)
+
+    parser.add_argument("--path", type=Path, required=True, help="Dataset root directory or JSONL file path.")
+    parser.add_argument("--database-name", type=str, default="M4_monolingual")
+    parser.add_argument(
+        "--model-name",
+        type=str,
+        default="FacebookAI/roberta-large",
+        help=(
+            "Model identifier for embeddings generation. Accepts either a Hugging Face "
+            "model hub name or a local path to a directory in Hugging Face format."
+        ),
+    )
+
+    parser.add_argument("--pooling", type=str, default="max", choices=("max", "average", "cls"))
+    parser.add_argument("--need-layer", type=int, nargs="+", default=[16, 17, 18, 19, 22, 23])
+
+    parser.add_argument("--adversarial", dest="adversarial", action="store_true")
+    parser.add_argument("--no-adversarial", dest="adversarial", action="store_false")
+    parser.set_defaults(adversarial=True)
+
+    parser.add_argument("--has-mix", dest="has_mix", action="store_true")
+    parser.add_argument("--no-has-mix", dest="has_mix", action="store_false")
+    parser.set_defaults(has_mix=False)
+
+    parser.add_argument("--savedir", type=Path, required=True, help="Output directory for the embedding database.")
+    parser.add_argument("--name", type=str, required=True, help="Filename (without extension) for the saved embeddings.")
+    parser.add_argument("--split", type=str, default="train", choices=("train", "test", "extra"))
+
+    return parser
+
+
+def generate_embeddings(args: argparse.Namespace) -> None:
+    if args.device_num > 1:
+        fabric = Fabric(accelerator="cuda", precision="bf16-mixed", devices=args.device_num, strategy="ddp")
+    else:
+        fabric = Fabric(accelerator="cuda", precision="bf16-mixed", devices=args.device_num)
+    fabric.launch()
+
+    model = TextEmbeddingModel(
+        args.model_name,
+        output_hidden_states=True,
+        infer=True,
+        use_pooling=args.pooling,
+    ).cuda()
+    tokenizer = model.tokenizer
+    model.eval()
+
+    path_str = str(args.path)
+    if "LLM_detect_data" in path_str:
+        now_data = load_data(args.split, include_adversarial=args.adversarial, fp=args.path)
+        now_data = now_data.to_dict(orient="records")
+        dataset = PassagesDataset(now_data)
+        dataloader = DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers)
+        dataloader = fabric.setup_dataloaders(dataloader)
+    elif path_str.endswith(".jsonl"):
+        dataset = SCLDataset([path_str], fabric, tokenizer, need_ids=True, adv_p=0)
+        dataloader = DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers)
+        dataloader = fabric.setup_dataloaders(dataloader, use_distributed_sampler=False)
+    else:
+        data_path = load_datapath(
+            path_str,
+            include_adversarial=args.adversarial,
+            dataset_name=args.database_name,
+        )[args.split]
+        dataset = SCLDataset(data_path, fabric, tokenizer, need_ids=True, adv_p=0, has_mix=args.has_mix)
+        dataloader = DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers)
+        dataloader = fabric.setup_dataloaders(dataloader, use_distributed_sampler=False)
+
+    model = fabric.setup(model)
+    classes = dataset.classes
+    train_ids, train_embeddings, train_labels = infer(dataloader, fabric, tokenizer, model, args)
+
+    torch.cuda.empty_cache()
+    if fabric.global_rank == 0:
+        args.savedir.mkdir(parents=True, exist_ok=True)
+        emb_dict = {
+            "embeddings": train_embeddings,
+            "labels": train_labels,
+            "ids": train_ids,
+            "classes": classes,
+        }
+        output_path = args.savedir / f"{args.name}.pt"
+        torch.save(emb_dict, output_path)
+        print(f"Saved embedding database to {output_path}")
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    args = parser.parse_args(argv)
+    generate_embeddings(args)
+
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["build_argument_parser", "generate_embeddings", "main"]

detree/cli/gen_tree.py

@@ -0,0 +1,86 @@
+"""Tree generation CLI utilities for DETree."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Iterable, Sequence, Set
+
+from detree.utils.dataset import load_datapath, model_alias_mapping
+
+
+def _str2bool(value: str) -> bool:
+    """Parse common textual boolean representations used by legacy scripts."""
+
+    if isinstance(value, bool):
+        return value
+    lowered = value.lower()
+    if lowered in {"true", "1", "yes", "y"}:
+        return True
+    if lowered in {"false", "0", "no", "n"}:
+        return False
+    raise argparse.ArgumentTypeError(f"Boolean value expected, got: {value}")
+
+
+def get_data_model(data_path: Iterable[Path], has_mix: bool = True) -> Set[str]:
+    """Collect all model identifiers present in the provided dataset paths."""
+
+    llm_name: Set[str] = set()
+    cnt = 0
+    for path in data_path:
+        print(f"reading {path}")
+        with path.open(mode="r", encoding="utf-8") as jsonl_file:
+            for line in jsonl_file:
+                now = json.loads(line)
+                if now["src"] not in model_alias_mapping:
+                    model_alias_mapping[now["src"]] = now["src"]
+                now["src"] = model_alias_mapping[now["src"]]
+                if not has_mix and "human" in now["src"] and now["src"] != "human":
+                    continue
+                if now["src"] not in llm_name:
+                    llm_name.add(now["src"])
+                cnt += 1
+    print(cnt)
+    return llm_name
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    """Create the argument parser for the tree generation CLI."""
+
+    parser = argparse.ArgumentParser(
+        description="Generate DETree-compatible tree definitions from dataset files.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--path", type=Path, default=Path("/opt/AI-text-Dataset"), help="Root directory of the dataset.")
+    parser.add_argument("--dataset_name", type=str, default="all", help="Dataset configuration name.")
+    parser.add_argument(
+        "--mode",
+        type=str,
+        choices=("train", "test", "extra"),
+        default="train",
+        help="Dataset split to consume.",
+    )
+    parser.add_argument("--tree_txt", type=Path, default=Path("output/Tree_RAID_pcl.txt"), help="Output tree definition path.")
+    parser.add_argument("--adversarial", type=_str2bool, default=True, help="Whether to include adversarial data splits.")
+    parser.add_argument("--has_mix", type=_str2bool, default=True, help="Whether to keep mixed human/model generations.")
+    return parser
+
+
+def main(args: argparse.Namespace) -> None:
+    """Entry point for building DETree-compatible tree structures."""
+
+    dataset_paths: Sequence[str] = load_datapath(args.path, args.adversarial, args.dataset_name)[args.mode]
+    print(f"data_path: {dataset_paths}")
+    llm_name = sorted(get_data_model((Path(p) for p in dataset_paths), args.has_mix))
+    root = len(llm_name)
+    args.tree_txt.parent.mkdir(parents=True, exist_ok=True)
+    with args.tree_txt.open("w", encoding="utf-8") as f:
+        for i, item in enumerate(llm_name):
+            f.write(f"{i} {root} {item}\n")
+        f.write(f"{root} -1 none\n")
+
+
+if __name__ == "__main__":
+    parser = build_argument_parser()
+    main(parser.parse_args())

detree/cli/hierarchical_clustering.py

@@ -0,0 +1,497 @@
+import argparse
+import random
+from pathlib import Path
+from typing import Iterable, Optional
+
+import matplotlib.cm as cm
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.cluster.hierarchy import dendrogram, linkage
+from scipy.spatial.distance import euclidean, squareform
+from sklearn.metrics import silhouette_score
+
+def read_similarity_matrix(file_path: Path):
+    with file_path.open('r', encoding='utf-8') as f:
+        lines = f.readlines()
+    names = lines[0].strip().split()
+    matrix = []
+
+    for line in lines[1:]:
+        row = line.strip().split()[1:]
+        matrix.append([float(x) for x in row])
+
+    similarity_matrix = np.array(matrix)
+    return names, similarity_matrix
+
+class TreeNode:
+    def __init__(self, name=None):
+
+        self.name = name
+        self.children = []
+        self.value = 0
+        self.split = True
+
+    def add_child(self, child):
+        self.children.append(child)
+        
+def build_tree(Z, names):
+    nodes = [TreeNode(name) for name in names]
+    for i, link in enumerate(Z):
+        node = TreeNode()
+        node.value = link[2]
+        node.add_child(int(link[0]))
+        node.add_child(int(link[1]))
+        nodes.append(node)
+    return nodes
+
+def find_best_thold(node_idx,nodes, distance_matrix,min_socre=0,max_socre=1):
+    node = nodes[node_idx]
+    threshold_range = np.linspace(min_socre * node.value, max_socre * node.value, 50)
+    silhouette_scores = []
+    all_n_clusters = []
+
+    for threshold in threshold_range:
+        labels,_ = gen_label_from_node(node_idx,nodes,threshold)
+        labels = sorted(labels,key=lambda x:x[1])
+        labels = [x[0] for x in labels]
+        n_clusters = len(np.unique(labels))
+        if n_clusters > 1 and n_clusters < len(distance_matrix):
+            score = silhouette_score(distance_matrix, labels, metric='precomputed')
+        else:
+            score = -1
+        silhouette_scores.append(score)
+        all_n_clusters.append(n_clusters)
+    best_threshold_idx = np.argmax(silhouette_scores)
+    best_threshold = threshold_range[best_threshold_idx]
+    best_score = silhouette_scores[best_threshold_idx]
+    return best_threshold, best_score
+
+def gen_label_from_node(node_idx,nodes,thd,now_label=0):
+    node = nodes[node_idx]
+    if len(node.children)==0:
+        return [(now_label,node_idx)],now_label
+    else:
+        if node.value>thd:
+            label_list = []
+            for child in node.children:
+                now_label_list,now_label = gen_label_from_node(child,nodes,thd,now_label)
+                now_label+=1
+                label_list+=now_label_list
+            return label_list,now_label
+        else:
+            label_list = []
+            for child in node.children:
+                now_label_list,now_label = gen_label_from_node(child,nodes,thd,now_label)
+                label_list+=now_label_list
+            return label_list,now_label
+
+def find_new_root(node_idx,nodes,thd):
+    node = nodes[node_idx]
+    if node.value<=thd:
+        return [node_idx]
+
+    new_root = []
+    for child in node.children:
+        new_root+=find_new_root(child,nodes,thd)
+    return new_root
+
+def get_leaf(node_idx,nodes):
+    node = nodes[node_idx]
+    if len(node.children)==0:
+        return [node_idx]
+
+    leaf_list = []
+    for child in node.children:
+        leaf_list+=get_leaf(child,nodes)
+    return leaf_list
+
+def merge_tree(node_idx,nodes,distance_matrix,deep=0,end_thd=0.25):
+    if len(nodes[node_idx].children)==0:
+        return
+    print(f"Node {node_idx}: Value: {nodes[node_idx].value}, Depth: {deep}")
+    if nodes[node_idx].value<=end_thd or deep>=5:
+        nodes[node_idx].children = get_leaf(node_idx,nodes)
+        nodes[node_idx].split = False
+        return
+    leaf_list = np.array(sorted(get_leaf(node_idx,nodes)))
+    new_distance_matrix = distance_matrix[leaf_list][:,leaf_list]
+    best_threshold, best_score = find_best_thold(node_idx, nodes, new_distance_matrix,min_socre=0)
+    if best_score==-1:
+        nodes[node_idx].children = get_leaf(node_idx,nodes)
+        return
+    new_root = find_new_root(node_idx,nodes,best_threshold)
+    nodes[node_idx].children = new_root
+
+    for child in new_root:
+        merge_tree(child,nodes,distance_matrix,deep=deep+1,end_thd=end_thd)
+
+def merge_dict(a,b):
+    for key in b.keys():
+        if key in a.keys():
+            a[key]+=b[key]
+        else:
+            a[key] = b[key]
+    return a
+
+def update_tree(node_idx, nodes, edge_list, fa=-1, deep=0):
+    node = nodes[node_idx]
+
+    if len(node.children)==0:
+        edge_list.append((fa,node_idx,[nodes[node_idx].name]))
+        return {deep:[[node_idx]]}
+
+    if node.split==False:
+        leafs = get_leaf(node_idx,nodes)
+        edge_list.append((fa,node_idx,[nodes[idx].name for idx in leafs]))
+        return {deep:[leafs]}
+
+    edge_list.append((fa,node_idx,[]))
+    new_tree = {}
+    for child in node.children:
+        new_tree = merge_dict(
+            new_tree,
+            update_tree(child, nodes, edge_list, node_idx, deep=deep+1),
+        )
+    if deep not in new_tree.keys():
+        new_tree[deep] = []
+    new_tree[deep].append(get_leaf(node_idx,nodes))
+
+    return new_tree
+
+def color_distance(c1, c2):
+    return euclidean(c1[:3], c2[:3])  # only consider the RGB components
+
+def ensure_color_diversity(colors, min_distance=0.2):
+    random.shuffle(colors)
+    for i in range(1, len(colors)):
+        if color_distance(colors[i], colors[i-1]) < min_distance:
+            for j in range(i + 1, len(colors)):
+                if color_distance(colors[i], colors[j]) > min_distance:
+                    colors[i], colors[j] = colors[j], colors[i]
+                    break
+    return colors
+
+
+def draw_table(new_tree, names, max_deep=3, save_path='fig/E/test.pdf'):
+    base_list = new_tree[0][0]
+    data = [base_list]
+    cmap = cm.get_cmap('tab20c', 2048)
+    cmap = [cmap(i) for i in range(2048)]
+    cmap = ensure_color_diversity(cmap)
+    cell_colours = [['#FFDDC1' for _ in base_list]]
+    color_start=0
+
+    for i in range(1,max_deep+1):
+        if i not in new_tree.keys():
+            print(f"Level {i} not in new_tree")
+            continue
+        data.append([names[base] for base in base_list])
+        color_list = []
+        for k,base in enumerate(base_list):
+            color_id = -1
+            for j in range(len(new_tree[i])):
+                if base in new_tree[i][j]:
+                    color_id = j
+                    break
+            if color_id==-1:
+                color_list.append(cell_colours[-1][k])
+            else:
+                color_list.append(cmap[color_start+color_id])
+        cell_colours.append(color_list)
+        color_start+=len(new_tree[i])
+
+    data = list(zip(*data))
+    cell_colours = list(zip(*cell_colours))            
+    columns = ['Node ID']+['Level {}'.format(i) for i in range(1,max_deep+1)]
+    plt.figure(figsize=(30, 40))
+    table = plt.table(cellText=data, colLabels=columns, loc='center', cellLoc='center',
+                colColours=['#f5f5f5']*len(columns),cellColours=cell_colours)
+    table.auto_set_column_width([0, 1]) 
+    plt.axis('off')
+    plt.savefig(save_path, format='pdf' ,bbox_inches='tight',pad_inches=0.01)
+
+def fix_asymmetry(matrix):
+    matrix = (matrix + matrix.T) / 2
+    return matrix
+
+def rename(edge):
+    cnt=0
+    reid={}
+    du={}
+    edge_dict={}
+    queue=[]
+    for i in range(len(edge)):
+        du[edge[i][0]]=du.get(edge[i][0],0)+1
+        edge_dict[edge[i][1]]=edge[i]
+        if edge[i][2] != []:
+            queue.append(edge[i][1])
+    while len(queue)>0:
+        now = queue.pop(0)
+        if now==-1:
+            reid[now]=-1
+            continue
+        if now not in reid.keys():
+            reid[now]=cnt
+            cnt+=1
+        now_edge = edge_dict[now]
+        du[now_edge[0]]-=1
+        if du[now_edge[0]]==0:
+            queue.append(now_edge[0])
+    new_edge = [(reid[x[0]],reid[x[1]],x[2]) for x in edge]
+    return new_edge
+
+def save_edge(edge,save_path):
+    with open(save_path,'w') as f:
+        for e in edge:
+            if e[2]:
+                name_str = ','.join(e[2])
+            else:
+                name_str = 'none'
+            f.write(f"{e[1]} {e[0]} {name_str}\n")
+
+def filter_class(names, similarity_matrix):
+    choose_idx = []
+    for i in range(len(names)):
+        if 'extend' not in names[i] and 'polish' not in names[i] and\
+              'translate' not in names[i] and 'paraphrase' not in names[i]:
+            if 'B' in names[i] or 'human' in names[i]:
+                choose_idx.append(i)
+            else:
+                if random.random()<0.3:
+                    choose_idx.append(i)
+        elif 'human' in names[i]:
+            if random.random()<0.3:
+                choose_idx.append(i)
+        elif random.random()<0.15:
+            choose_idx.append(i)
+    new_names = [names[i] for i in choose_idx]
+    choose_idx  = np.array(choose_idx)
+    new_similarity_matrix = similarity_matrix[choose_idx][:,choose_idx]
+    return new_names, new_similarity_matrix
+
+def filter(names, similarity_matrix,filter_human=False,filter_llm=False,filter_mix=False):
+    choose_idx = []
+    for i in range(len(names)):
+        if names[i] == 'human' and filter_human:
+            continue
+        if filter_llm and 'human' not in names[i]:
+            continue
+        if filter_mix and 'human' in names[i] and names[i]!='human':
+            continue
+        choose_idx.append(i)
+    new_names = [names[i] for i in choose_idx]
+    choose_idx  = np.array(choose_idx)
+    new_similarity_matrix = similarity_matrix[choose_idx][:,choose_idx]
+    return new_names, new_similarity_matrix
+
+def reid_tree_dict(tree_dict, nodes, names):
+    name_to_index = {name: idx for idx, name in enumerate(names)}
+    for deep,values in tree_dict.items():
+        rename_now = []
+        # print(values,len(values))
+        for list_ in values:
+            now_list = []
+            for idx in list_:
+                name = nodes[idx].name
+                if name not in name_to_index:
+                    name_to_index[name] = len(names)
+                    names.append(name)
+                name_idx = name_to_index[name]
+                now_list.append(name_idx)
+            rename_now.append(now_list)
+        tree_dict[deep] = rename_now
+    return tree_dict
+
+def gen_tree(similarity_matrix,names,opt):
+    distance_matrix = 1 - similarity_matrix
+    np.fill_diagonal(distance_matrix, 0)
+    condensed_distance_matrix = squareform(distance_matrix)
+    Z = linkage(condensed_distance_matrix, method='weighted')  # alternative methods include 'single', 'complete', or 'ward'
+    if opt.save_drg:
+        plt.figure(figsize=(30, 47))
+        dendrogram(Z, labels=names, orientation='right',leaf_font_size=16)  # rotate the dendrogram so the root is on the right
+        plt.savefig(opt.dendrogram_path, format='pdf' ,bbox_inches='tight')
+    nodes = build_tree(Z, names)
+    merge_tree(len(nodes)-1,nodes,distance_matrix,end_thd=opt.end_score)
+
+    return nodes
+
+def chage_tree_priori1(nodes):
+    human_node = TreeNode(name='human')
+    root = TreeNode()
+    root.add_child(len(nodes))
+    root.add_child(len(nodes)-1)
+    nodes.append(human_node)
+    nodes.append(root)
+    return nodes
+
+def chage_tree_priori2(human_nodes,llm_nodes):
+    root = TreeNode()
+    root.add_child(len(human_nodes)-1)
+    root.add_child(len(human_nodes)+len(llm_nodes)-1)
+    for i in range(len(llm_nodes)):
+        llm_nodes[i].children = [len(human_nodes)+x for x in llm_nodes[i].children]
+    nodes = human_nodes+llm_nodes
+    nodes.append(root)
+    return nodes
+
+def chage_tree_priori3(co_nodes,llm_nodes):
+    human_node = TreeNode(name='human')
+    root = TreeNode()
+    root.add_child(len(co_nodes)+len(llm_nodes))
+    root.add_child(len(co_nodes)-1)
+    root.add_child(len(co_nodes)+len(llm_nodes)-1)
+    for i in range(len(llm_nodes)):
+        llm_nodes[i].children = [len(co_nodes)+x for x in llm_nodes[i].children]
+    nodes = co_nodes+llm_nodes
+    nodes.append(human_node)
+    nodes.append(root)
+    return nodes
+
+def randmo_filter(names, similarity_matrix):
+    choose_idx = []
+    for i in range(len(names)):
+        if 'human' in names[i]:
+                choose_idx.append(i)
+        elif 'fair' in names[i] or 'pplm' in names[i] or 'gpt2-pytorch' in names[i] or ' transfo' in names[i]  or 'ctrl' in names[i]:
+                continue
+        elif 'xlnet' in names[i] or 'grover' in names[i]:
+            if random.random()<0.07:
+                choose_idx.append(i)
+        elif random.random()<0.22:
+                choose_idx.append(i)
+    new_names = []
+    for i in choose_idx:
+        if names[i].startswith('7B') or names[i].startswith('13B') or names[i].startswith('30B') or names[i].startswith('65B'):
+            new_names.append('LLaMA_'+names[i])
+        else:
+            new_names.append(names[i])
+    choose_idx  = np.array(choose_idx)
+    new_similarity_matrix = similarity_matrix[choose_idx][:,choose_idx]
+    return new_names, new_similarity_matrix
+
+def ishuman(name):
+    return ('human' in name)
+def ismachine(name):
+    return ('machine' in name or 'rephrase' in name)
+
+def get_llm(x):
+    if 'gpt-3.5-turbo' in x:
+        return 'gpt-3.5-turbo'
+    elif 'gpt-4o' in x:
+        return 'gpt-4o'
+    elif 'llama-3.3-70b' in x:
+        return 'llama-3.3-70b'
+    elif 'gemini-1.5-pro' in x:
+        return 'gemini-1.5-pro'
+    elif 'claude-3-5-sonnet' in x:
+        return 'claude-3-5-sonnet'
+    elif 'qwen2.5-72b' in x:
+        return 'qwen2.5-72b'
+    else:
+        raise ValueError(f"Invalid class name: {x}")
+    
+def get_name(name):
+    name = name.split('_')
+    assert len(name) == 2
+    if ishuman(name[0]):
+        if name[1]=='humanize:human' or name[1]=='human':
+            return 'human'
+        elif name[1]=='humanize:tool':
+            return 'human_humanize_tool'
+        else:
+            llm_name = get_llm(name[1])
+            return f'human_rephrase_{llm_name}'
+    elif ismachine(name[0]):
+        llm_name = get_llm(name[0])
+        if name[1]=='humanize:human' or name[1]=='human':
+            return f'{llm_name}_humanize_human'
+        elif name[1]=='humanize:tool':
+            return f'{llm_name}_humanize_tool'
+        elif 'humanize:' in name[1]:
+            llm_name2 = get_llm(name[1])
+            return f'{llm_name}_humanize_{llm_name2}'
+        else:
+            return llm_name
+
+def clear_names(names):
+    new_names = []
+    for name in names:
+        new_names.append(get_name(name))
+    return new_names
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Construct the HAT tree from a similarity matrix.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument('--file-path', type=Path, required=True, help='Input similarity matrix text file.')
+    parser.add_argument('--priori',type=int,default=1,choices=[0,1,2,3])
+    parser.add_argument('--save-txt-path', type=Path, required=True, help='Destination path for the tree definition.')
+    parser.add_argument('--save-table-path', type=Path, required=True, help='Destination path for the visualised table.')
+    parser.add_argument('--dendrogram-path', type=Path, default=None, help='Optional path for the dendrogram PDF when saved.')
+    parser.add_argument('--save-drg', action='store_true', help='Persist the dendrogram PDF alongside the tree.')
+    parser.add_argument('--no-save-drg', dest='save_drg', action='store_false')
+    parser.set_defaults(save_drg=True)
+    parser.add_argument('--save-max-dep', type=int, default=5)
+    parser.add_argument('--end-score', type=float, default=0.1)
+    parser.add_argument('--randmo-filter', action='store_true', help='Randomly subsample similarity entries.')
+    return parser
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    opt = parser.parse_args(argv)
+
+    names, similarity_matrix = read_similarity_matrix(opt.file_path)
+    if opt.save_drg:
+        if opt.dendrogram_path is None:
+            opt.dendrogram_path = opt.save_table_path.with_name(
+                f"{opt.save_table_path.stem}_dendrogram.pdf"
+            )
+        opt.dendrogram_path.parent.mkdir(parents=True, exist_ok=True)
+    else:
+        opt.dendrogram_path = None
+    similarity_matrix = fix_asymmetry(similarity_matrix)
+    if opt.randmo_filter:
+        names, similarity_matrix = randmo_filter(names, similarity_matrix)
+    # names = clear_names(names)
+    if opt.priori==1:
+        llm_names, llm_similarity_matrix = filter(names, similarity_matrix,filter_human=True)
+        nodes = gen_tree(llm_similarity_matrix,llm_names,opt)
+        nodes = chage_tree_priori1(nodes)
+
+    elif opt.priori==2:
+        human_names, human_similarity_matrix = filter(names, similarity_matrix,filter_llm=True)
+        human_nodes = gen_tree(human_similarity_matrix,human_names,opt)
+        llm_names, llm_similarity_matrix = filter(names, similarity_matrix,filter_human=True,filter_mix=True)
+        llm_nodes = gen_tree(llm_similarity_matrix,llm_names,opt)
+        nodes = chage_tree_priori2(human_nodes,llm_nodes)
+    
+    elif opt.priori==3:
+        co_names, co_similarity_matrix = filter(names, similarity_matrix,filter_llm=True,filter_human=True)
+        co_nodes = gen_tree(co_similarity_matrix,co_names,opt)
+        llm_names, llm_similarity_matrix = filter(names, similarity_matrix,filter_human=True,filter_mix=True)
+        llm_nodes = gen_tree(llm_similarity_matrix,llm_names,opt)
+        nodes = chage_tree_priori3(co_nodes,llm_nodes)
+    
+    elif opt.priori==0:
+        nodes = gen_tree(similarity_matrix,names,opt)
+    else:
+        raise ValueError("Invalid value for --priori. Choose from 0, 1, 2, or 3.")
+
+    edge=[]
+    tree_dict = update_tree(len(nodes)-1, nodes, edge)
+    edge = rename(edge)
+    opt.save_txt_path.parent.mkdir(parents=True, exist_ok=True)
+    opt.save_table_path.parent.mkdir(parents=True, exist_ok=True)
+    save_edge(edge,opt.save_txt_path)
+    tree_dict = reid_tree_dict(tree_dict, nodes, names)
+    draw_table(tree_dict, names, max_deep=opt.save_max_dep, save_path=opt.save_table_path)
+
+
+if __name__ == "__main__":
+    main()
+
+
+__all__ = ["build_argument_parser", "main", "read_similarity_matrix", "gen_tree"]

detree/cli/merge_lora.py

@@ -0,0 +1,52 @@
+"""Merge LoRA adapters into base models."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Iterable, Optional
+
+from peft import PeftModel
+from transformers import AutoModel, AutoTokenizer
+
+
+def merge_lora_adapter(base_model: str, adapter_path: Path, output_dir: Path, safe_serialization: bool = True) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    model = AutoModel.from_pretrained(base_model, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(base_model)
+
+    peft_model = PeftModel.from_pretrained(model, str(adapter_path))
+    merged_model = peft_model.merge_and_unload()
+    merged_model.save_pretrained(output_dir, safe_serialization=safe_serialization)
+    tokenizer.save_pretrained(output_dir)
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Merge a LoRA adapter into its base Hugging Face model.")
+    parser.add_argument("--base-model", type=str, required=True, help="Base model name or path.")
+    parser.add_argument("--adapter-path", type=Path, required=True, help="Directory containing the LoRA adapter weights.")
+    parser.add_argument("--output-dir", type=Path, required=True, help="Directory to store the merged model.")
+    parser.add_argument(
+        "--no-safe-serialization",
+        action="store_true",
+        help="Disable safetensors when saving the merged model.",
+    )
+    return parser
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    args = parser.parse_args(argv)
+    merge_lora_adapter(
+        args.base_model,
+        args.adapter_path,
+        args.output_dir,
+        safe_serialization=not args.no_safe_serialization,
+    )
+
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["build_argument_parser", "merge_lora_adapter", "main"]

detree/cli/similarity_matrix.py

@@ -0,0 +1,77 @@
+"""Compute similarity matrices from embedding databases."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Iterable, Optional
+
+import matplotlib.pyplot as plt
+import torch
+
+
+def gen_data(dict_data):
+    embeddings = dict_data["embeddings"]
+    labels = dict_data["labels"]
+    ids = dict_data["ids"]
+    classes = dict_data["classes"]
+    return embeddings, labels, ids, classes
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Generate class similarity matrices for DETree.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--database", type=Path, required=True, help="Path to the embedding database (.pt).")
+    parser.add_argument("--output-dir", type=Path, required=True, help="Directory to store the similarity outputs.")
+    parser.add_argument("--layers", type=int, nargs="*", default=None, help="Specific layers to export. Defaults to all.")
+    return parser
+
+
+def compute_similarity(database: Path, output_dir: Path, layers: Optional[Iterable[int]]) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    data_emb, data_labels, data_ids, data_classes = gen_data(torch.load(database))
+
+    if layers is None:
+        layers = list(data_emb.keys())
+
+    for layer in layers:
+        center = []
+        for item in data_classes:
+            index = data_classes.index(item)
+            now_emb = data_emb[layer][data_labels == index]
+            center.append(torch.mean(now_emb, dim=0))
+        center = torch.stack(center)
+        similarity = center @ center.T
+        similarity_np = similarity.cpu().float().numpy()
+
+        txt_path = output_dir / f"similarity_layer_{layer}.txt"
+        with txt_path.open("w", encoding="utf-8") as f:
+            f.write(" ".join(data_classes) + "\n")
+            for i, class_name in enumerate(data_classes):
+                row = " ".join(f"{similarity_np[i, j]:.4f}" for j in range(len(data_classes)))
+                f.write(f"{class_name} {row}\n")
+
+        plt.figure(figsize=(30, 30))
+        plt.imshow(similarity_np, cmap="viridis")
+        plt.colorbar()
+        plt.xticks(range(len(data_classes)), data_classes, rotation=45, fontsize=12)
+        plt.yticks(range(len(data_classes)), data_classes, fontsize=12)
+        plt.title(f"Similarity Matrix (layer {layer})", fontsize=20)
+        fig_path = output_dir / f"similarity_layer_{layer}.png"
+        plt.savefig(fig_path, dpi=300, bbox_inches="tight")
+        plt.close()
+        print(f"Saved similarity matrix for layer {layer} to {txt_path} and {fig_path}")
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    args = parser.parse_args(argv)
+    compute_similarity(args.database, args.output_dir, args.layers)
+
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["build_argument_parser", "compute_similarity", "main"]

detree/cli/test_database_score_knn.py

@@ -0,0 +1,247 @@
+"""kNN evaluation using pre-computed embedding databases."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from multiprocessing import Pool, cpu_count
+from pathlib import Path
+from typing import Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from lightning import Fabric
+from torch.nn.functional import softmax as F_softmax
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from detree.model.text_embedding import TextEmbeddingModel
+from detree.utils.index import Indexer
+from detree.utils.utils import evaluate_metrics
+
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")
+
+
+def load_jsonl(file_path: Path) -> List[dict]:
+    out = []
+    with file_path.open(mode="r", encoding="utf-8") as jsonl_file:
+        for line in jsonl_file:
+            item = json.loads(line)
+            out.append(item)
+    print(f"Loaded {len(out)} examples from {file_path}")
+    return out
+
+
+def gen_data(dict_data):
+    embeddings = dict_data["embeddings"]
+    labels = dict_data["labels"]
+    ids = dict_data["ids"]
+    classes = dict_data["classes"]
+    return embeddings, labels, ids, classes
+
+
+class PassagesDataset(Dataset):
+    def __init__(self, data: Sequence[dict]):
+        self.passages = list(data)
+
+    def __len__(self) -> int:
+        return len(self.passages)
+
+    def __getitem__(self, idx: int):
+        data_now = self.passages[idx]
+        text = data_now["text"]
+        label = data_now["label"]
+        ids = data_now["id"]
+        return text, int(label), int(ids)
+
+
+def infer(passages_dataloader, fabric, tokenizer, model, need_layers: Sequence[int], max_length: int = 512):
+    if fabric.global_rank == 0:
+        passages_dataloader = tqdm(passages_dataloader)
+        all_ids: List[int] = []
+        all_embeddings: List[torch.Tensor] = []
+        all_labels: List[int] = []
+    with torch.no_grad():
+        for batch in passages_dataloader:
+            text, label, ids = batch
+            encoded_batch = tokenizer.batch_encode_plus(
+                text,
+                return_tensors="pt",
+                max_length=max_length,
+                padding="max_length",
+                truncation=True,
+            )
+            encoded_batch = {k: v.cuda() for k, v in encoded_batch.items()}
+            embeddings = model(encoded_batch, hidden_states=True)
+            embeddings = fabric.all_gather(embeddings).view(-1, embeddings.size(-2), embeddings.size(-1))
+            label = fabric.all_gather(label).view(-1)
+            ids = fabric.all_gather(ids).view(-1)
+            if fabric.global_rank == 0:
+                all_embeddings.append(embeddings.cpu())
+                all_ids.extend(ids.cpu().tolist())
+                all_labels.extend(label.cpu().tolist())
+    if fabric.global_rank == 0:
+        embeddings_tensor = torch.cat(all_embeddings, dim=0)
+        embeddings_tensor = F.normalize(embeddings_tensor, dim=-1).permute(1, 0, 2).numpy()
+        embeddings_tensor = {layer: embeddings_tensor[layer] for layer in need_layers}
+        return all_ids, embeddings_tensor, all_labels
+    return [], [], []
+
+
+def dict2str(metrics: dict) -> str:
+    out_str = ""
+    if "layer" in metrics:
+        out_str += f"layer:{metrics['layer']} "
+    if "k" in metrics:
+        out_str += f"k:{metrics['k']} "
+    for key, value in metrics.items():
+        if key not in {"layer", "k"}:
+            out_str += f"{key}:{value} "
+    return out_str.strip()
+
+
+def process_element(args: Tuple[Sequence[int], Sequence[float], Sequence[int], float]):
+    ids, scores, labels, temperature = args
+    now_score = torch.zeros(2)
+    sorted_indices = np.argsort(scores)[::-1]
+    element_preds = {}
+
+    for k, idx in enumerate(sorted_indices):
+        label = labels[idx]
+        now_score[label] += scores[idx] * temperature
+        prob = F_softmax(now_score, dim=-1)[1].item()
+        element_preds[k + 1] = prob
+
+    return element_preds
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Evaluate DETree with a precomputed embedding database.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--device-num", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-workers", type=int, default=8)
+    parser.add_argument("--max-length", type=int, default=512)
+
+    parser.add_argument("--database-path", type=Path, required=True, help="Path to the saved embedding database (.pt).")
+    parser.add_argument("--test-dataset-path", type=Path, required=True, help="Evaluation JSONL file.")
+    parser.add_argument("--model-name-or-path", type=str, required=True)
+    parser.add_argument("--temperature", type=float, default=0.05)
+
+    parser.add_argument("--max-k", type=int, default=51, dest="max_K")
+    parser.add_argument("--pooling", type=str, default="max", choices=("max", "average", "cls"))
+
+    parser.add_argument("--embedding-dim", type=int, default=1024)
+    parser.add_argument("--pool-workers", type=int, default=min(32, cpu_count()))
+    parser.add_argument("--log-file", type=Path, default=Path("runs/val.txt"))
+
+    return parser
+
+
+def evaluate(args: argparse.Namespace) -> None:
+    if args.device_num > 1:
+        fabric = Fabric(accelerator="cuda", devices=args.device_num, strategy="ddp", precision="bf16-mixed")
+    else:
+        fabric = Fabric(accelerator="cuda", devices=args.device_num, precision="bf16-mixed")
+    fabric.launch()
+
+    model = TextEmbeddingModel(
+        args.model_name_or_path,
+        output_hidden_states=True,
+        infer=True,
+        use_pooling=args.pooling,
+    ).cuda()
+    tokenizer = model.tokenizer
+    model.eval()
+
+    if fabric.global_rank == 0:
+        db_embeddings, db_labels, db_ids, classes = gen_data(torch.load(args.database_path))
+        need_layers = list(db_embeddings.keys())
+    else:
+        db_embeddings = db_labels = db_ids = classes = None
+        need_layers = []
+    need_layers = fabric.broadcast(need_layers)
+
+    test_database = load_jsonl(args.test_dataset_path)
+    test_dataset = PassagesDataset(test_database)
+    test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=False)
+    test_dataloader = fabric.setup_dataloaders(test_dataloader)
+    model = fabric.setup(model)
+    test_ids, test_embeddings, test_labels = infer(test_dataloader, fabric, tokenizer, model, need_layers, args.max_length)
+
+    torch.cuda.empty_cache()
+    if fabric.global_rank != 0:
+        return
+
+    test_labels = [int(label) for label in test_labels]
+    index = Indexer(args.embedding_dim)
+    human_idx = classes.index("human")
+
+    all_details = []
+    with Pool(processes=args.pool_workers) as pool:
+        for layer in need_layers:
+            now_best_metrics = None
+            label_dict = {}
+            train_embeddings = db_embeddings[layer].float().numpy()
+            if isinstance(db_labels, dict):
+                train_labels = db_labels[layer].tolist()
+                train_ids = db_ids[layer].tolist()
+            else:
+                train_labels = db_labels.tolist()
+                train_ids = db_ids.tolist()
+
+            for i in range(len(train_ids)):
+                label_dict[int(train_ids[i])] = int(train_labels[i] == human_idx)
+            index.label_dict = label_dict
+            index.reset()
+            index.index_data(train_ids, train_embeddings)
+            preds = {k: [] for k in range(1, args.max_K + 1)}
+            top_ids_and_scores = index.search_knn(test_embeddings[layer], args.max_K, index_batch_size=128)
+
+            args_list = [
+                (ids, scores, labels, args.temperature)
+                for ids, scores, labels in top_ids_and_scores
+            ]
+            for result in tqdm(pool.imap(process_element, args_list), total=len(args_list)):
+                for k, value in result.items():
+                    preds[k].append(value)
+
+            for k in range(1, args.max_K + 1):
+                metric = evaluate_metrics(test_labels, preds[k], threshold_param=-1)
+                if now_best_metrics is None or now_best_metrics["auroc"] < metric["auroc"]:
+                    now_best_metrics = metric
+                    now_best_metrics["k"] = k
+                    now_best_metrics["layer"] = layer
+
+            if now_best_metrics:
+                print(dict2str(now_best_metrics))
+                all_details.append(now_best_metrics)
+
+    if not all_details:
+        return
+
+    max_ids = max(range(len(all_details)), key=lambda idx: all_details[idx]["auroc"])
+    best_metrics = all_details[max_ids]
+    print("Best " + dict2str(best_metrics))
+    args.log_file.parent.mkdir(parents=True, exist_ok=True)
+    with args.log_file.open("a+", encoding="utf-8") as fp:
+        fp.write(f"test model:{args.model_name_or_path} mode:{args.test_dataset_path} database_path:{args.database_path}\n")
+        fp.write(f"Last {dict2str(all_details[-1])}\n")
+        fp.write(f"Best {dict2str(best_metrics)}\n")
+        fp.write("------------------------------------------\n")
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    args = parser.parse_args(argv)
+    evaluate(args)
+
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["build_argument_parser", "evaluate", "main"]

detree/cli/test_score_knn.py

@@ -0,0 +1,267 @@
+"""kNN evaluation against raw text datasets."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from multiprocessing import Pool, cpu_count
+from pathlib import Path
+from typing import Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from lightning import Fabric
+from torch.nn.functional import softmax as F_softmax
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from detree.model.text_embedding import TextEmbeddingModel
+from detree.utils.index import Indexer
+from detree.utils.utils import evaluate_metrics
+
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")
+
+
+def load_jsonl(file_path: Path) -> List[dict]:
+    out = []
+    with file_path.open(mode="r", encoding="utf-8") as jsonl_file:
+        for line in jsonl_file:
+            item = json.loads(line)
+            out.append(item)
+    print(f"Loaded {len(out)} examples from {file_path}")
+    return out
+
+
+class PassagesDataset(Dataset):
+    def __init__(self, data: Sequence[dict]):
+        self.passages = list(data)
+
+    def __len__(self) -> int:
+        return len(self.passages)
+
+    def __getitem__(self, idx: int):
+        data_now = self.passages[idx]
+        text = data_now["text"]
+        label = data_now["label"]
+        ids = data_now["id"]
+        return text, int(label), int(ids)
+
+
+def infer(passages_dataloader, fabric, tokenizer, model, max_length: int = 512):
+    if fabric.global_rank == 0:
+        passages_dataloader = tqdm(passages_dataloader)
+        all_ids: List[int] = []
+        all_embeddings: List[torch.Tensor] = []
+        all_labels: List[int] = []
+    with torch.no_grad():
+        for batch in passages_dataloader:
+            text, label, ids = batch
+            encoded_batch = tokenizer.batch_encode_plus(
+                text,
+                return_tensors="pt",
+                max_length=max_length,
+                padding="max_length",
+                truncation=True,
+            )
+            encoded_batch = {k: v.cuda() for k, v in encoded_batch.items()}
+            embeddings = model(encoded_batch, hidden_states=True)
+            embeddings = fabric.all_gather(embeddings).view(-1, embeddings.size(-2), embeddings.size(-1))
+            label = fabric.all_gather(label).view(-1)
+            ids = fabric.all_gather(ids).view(-1)
+            if fabric.global_rank == 0:
+                all_embeddings.append(embeddings.cpu())
+                all_ids.extend(ids.cpu().tolist())
+                all_labels.extend(label.cpu().tolist())
+    if fabric.global_rank == 0:
+        embeddings_tensor = torch.cat(all_embeddings, dim=0)
+        embeddings_tensor = F.normalize(embeddings_tensor, dim=-1).permute(1, 0, 2)
+        return all_ids, embeddings_tensor.numpy(), all_labels
+    return [], [], []
+
+
+def save_pt(train_embeddings, all_labels, train_ids, args, best_layer):
+    save_layer = [best_layer, train_embeddings.shape[0] - 1]
+    all_embeddings = {i: torch.tensor(train_embeddings[i]) for i in save_layer}
+    emb_dict = {
+        "embeddings": all_embeddings,
+        "labels": torch.tensor(all_labels),
+        "ids": torch.tensor(train_ids),
+        "classes": ["llm", "human"],
+    }
+    args.savedir.mkdir(parents=True, exist_ok=True)
+    output_path = args.savedir / f"{args.name}.pt"
+    torch.save(emb_dict, output_path)
+    print(f"Saved embedding snapshot to {output_path}")
+
+
+def dict2str(metrics: dict) -> str:
+    out_str = ""
+    if "layer" in metrics:
+        out_str += f"layer:{metrics['layer']} "
+    if "k" in metrics:
+        out_str += f"k:{metrics['k']} "
+    for key, value in metrics.items():
+        if key not in {"layer", "k"}:
+            out_str += f"{key}:{value} "
+    return out_str.strip()
+
+
+def process_element(args: Tuple[Sequence[int], Sequence[float], Sequence[int], float]):
+    ids, scores, labels, temperature = args
+    now_score = torch.zeros(2)
+    sorted_indices = np.argsort(scores)[::-1]
+    element_preds = {}
+
+    for k, idx in enumerate(sorted_indices):
+        label = labels[idx]
+        now_score[label] += scores[idx] * temperature
+        prob = F_softmax(now_score, dim=-1)[1].item()
+        element_preds[k + 1] = prob
+
+    return element_preds
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Evaluate DETree checkpoints using a kNN classifier over hidden states.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--device-num", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-workers", type=int, default=8)
+    parser.add_argument("--max-length", type=int, default=512)
+
+    parser.add_argument("--database-path", type=Path, required=True, help="Training set JSONL file.")
+    parser.add_argument("--test-dataset-path", type=Path, required=True, help="Evaluation set JSONL file.")
+    parser.add_argument(
+        "--model-name-or-path",
+        type=str,
+        required=True,
+        help="Model identifier from Hugging Face or local path to a merged checkpoint.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.05)
+
+    parser.add_argument("--max-k", type=int, default=50, dest="max_K", help="Maximum k to evaluate for kNN.")
+    parser.add_argument("--min-layer", type=int, default=15, help="Minimum hidden layer index to evaluate.")
+    parser.add_argument("--pooling", type=str, default="max", choices=("max", "average", "cls"))
+
+    parser.add_argument("--embedding-dim", type=int, default=1024)
+    parser.add_argument("--n-subquantizers", type=int, default=1)
+    parser.add_argument("--n-bits", type=int, default=8)
+
+    parser.add_argument("--savedir", type=Path, default=Path("runs"))
+    parser.add_argument("--name", type=str, default="database_knn_eval")
+    parser.add_argument("--pool-workers", type=int, default=min(32, cpu_count()))
+    parser.add_argument("--save-embeddings", action="store_true", help="Persist embeddings for the best-performing layer.")
+    parser.add_argument("--log-file", type=Path, default=Path("runs/val.txt"))
+
+    return parser
+
+
+def evaluate(args: argparse.Namespace) -> None:
+    if args.device_num > 1:
+        fabric = Fabric(accelerator="cuda", devices=args.device_num, strategy="ddp", precision="bf16-mixed")
+    else:
+        fabric = Fabric(accelerator="cuda", devices=args.device_num, precision="bf16-mixed")
+    fabric.launch()
+
+    model = TextEmbeddingModel(
+        args.model_name_or_path,
+        output_hidden_states=True,
+        infer=True,
+        use_pooling=args.pooling,
+    ).cuda()
+    tokenizer = model.tokenizer
+    model.eval()
+
+    database = load_jsonl(args.database_path)
+    test_database = load_jsonl(args.test_dataset_path)
+
+    passages_dataset = PassagesDataset(database)
+    test_dataset = PassagesDataset(test_database)
+
+    passages_dataloader = DataLoader(
+        passages_dataset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True
+    )
+    test_dataloader = DataLoader(
+        test_dataset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=False
+    )
+
+    passages_dataloader, test_dataloader = fabric.setup_dataloaders(passages_dataloader, test_dataloader)
+    model = fabric.setup(model)
+
+    train_ids, train_embeddings, train_labels = infer(passages_dataloader, fabric, tokenizer, model, args.max_length)
+    test_ids, test_embeddings, test_labels = infer(test_dataloader, fabric, tokenizer, model, args.max_length)
+
+    torch.cuda.empty_cache()
+    if fabric.global_rank != 0:
+        return
+
+    layer_num = train_embeddings.shape[0]
+    test_labels = [int(label) for label in test_labels]
+
+    label_dict = {train_ids[i]: int(train_labels[i]) for i in range(len(train_ids))}
+
+    all_details = []
+    index = Indexer(args.embedding_dim, args.n_subquantizers, args.n_bits)
+    index.label_dict = label_dict
+
+    with Pool(processes=args.pool_workers) as pool:
+        for i in range(args.min_layer, layer_num):
+            now_best_metrics = None
+            index.reset()
+            index.index_data(train_ids, train_embeddings[i])
+            preds = {k: [] for k in range(1, args.max_K + 1)}
+            top_ids_and_scores = index.search_knn(test_embeddings[i], args.max_K, index_batch_size=128)
+
+            args_list = [
+                (ids, scores, labels, args.temperature)
+                for ids, scores, labels in top_ids_and_scores
+            ]
+            for result in tqdm(pool.imap(process_element, args_list), total=len(args_list)):
+                for k, value in result.items():
+                    preds[k].append(value)
+
+            for k in range(2, args.max_K + 1):
+                metric = evaluate_metrics(test_labels, preds[k], threshold_param=-1)
+                if now_best_metrics is None or now_best_metrics["auroc"] < metric["auroc"]:
+                    now_best_metrics = metric
+                    now_best_metrics["k"] = k
+                    now_best_metrics["layer"] = i
+
+            if now_best_metrics:
+                print(dict2str(now_best_metrics))
+                all_details.append(now_best_metrics)
+
+    if not all_details:
+        return
+
+    max_ids = max(range(len(all_details)), key=lambda idx: all_details[idx]["auroc"])
+    best_metrics = all_details[max_ids]
+
+    if args.save_embeddings:
+        save_pt(train_embeddings, train_labels, train_ids, args, best_metrics["layer"])
+
+    print("Best " + dict2str(best_metrics))
+    args.log_file.parent.mkdir(parents=True, exist_ok=True)
+    with args.log_file.open("a+", encoding="utf-8") as fp:
+        fp.write(
+            f"test model:{args.model_name_or_path} database_path:{args.database_path} mode:{args.test_dataset_path}\n"
+        )
+        fp.write(f"Last {dict2str(all_details[-1])}\n")
+        fp.write(f"Best {dict2str(best_metrics)}\n")
+        fp.write("------------------------------------------\n")
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    args = parser.parse_args(argv)
+    evaluate(args)
+
+
+if __name__ == "__main__":
+    main()
+
+__all__ = ["build_argument_parser", "evaluate", "main"]

detree/cli/train.py

@@ -0,0 +1,313 @@
+"""Training CLI for DETree."""
+
+from __future__ import annotations
+
+import argparse
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Optional
+
+import torch
+import torch.nn.functional as F  # noqa: F401  # retained for backward compat with downstream imports
+import torch.optim as optim
+import yaml
+from lightning import Fabric
+from lightning.fabric.strategies import DeepSpeedStrategy, DDPStrategy
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import default_collate
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+from detree.model.simclr import SimCLR_Tree
+from detree.utils.dataset import SCLDataset, load_datapath
+
+
+@dataclass
+class ExperimentPaths:
+    """Utility container describing where to store experiment artefacts."""
+
+    root: Path
+    runs: Path
+
+
+def _build_collate_fn(tokenizer, max_length: int):
+    def collate_fn(batch: Iterable):
+        text, label, write_model = default_collate(batch)
+        encoded_batch = tokenizer.batch_encode_plus(
+            text,
+            return_tensors="pt",
+            max_length=max_length,
+            padding=True,
+            truncation=True,
+        )
+        return encoded_batch, label, write_model
+
+    return collate_fn
+
+
+def _prepare_output_dir(
+    output_dir: Path, experiment_name: str, resume: bool, *, create_dirs: bool = True
+) -> ExperimentPaths:
+    output_dir = output_dir.expanduser().resolve()
+
+    candidate = output_dir / experiment_name
+    if candidate.exists() and not resume:
+        suffix = 0
+        while (output_dir / f"{experiment_name}_v{suffix}").exists():
+            suffix += 1
+        candidate = output_dir / f"{experiment_name}_v{suffix}"
+
+    runs_dir = candidate / "runs"
+    if create_dirs:
+        candidate.mkdir(parents=True, exist_ok=True)
+        runs_dir.mkdir(parents=True, exist_ok=True)
+
+    return ExperimentPaths(root=candidate, runs=runs_dir)
+
+
+def build_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Train DETree using the hierarchical contrastive objective",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--model-name", type=str, default="FacebookAI/roberta-large", help="Backbone encoder identifier.")
+    parser.add_argument("--device-num", type=int, default=1, help="Number of CUDA devices to use.")
+    parser.add_argument("--path", type=Path, required=True, help="Root directory of the dataset.")
+    parser.add_argument("--dataset-name", type=str, default="all", help="Dataset configuration name.")
+    parser.add_argument(
+        "--dataset", type=str, default="train", choices=("train", "test", "extra"), help="Dataset split to consume."
+    )
+    parser.add_argument("--tree-txt", type=Path, required=True, help="Tree definition file as produced by the HAT pipeline.")
+    parser.add_argument("--output-dir", type=Path, default=Path("runs"), help="Directory where experiment folders are saved.")
+    parser.add_argument("--experiment-name", type=str, default="detree_experiment", help="Base name for the run directory.")
+    parser.add_argument("--resume", action="store_true", help="Reuse the given experiment directory if it already exists.")
+
+    parser.add_argument("--projection-size", type=int, default=1024)
+    parser.add_argument("--temperature", type=float, default=0.07)
+    parser.add_argument("--num-workers", type=int, default=8)
+    parser.add_argument("--per-gpu-batch-size", type=int, default=64)
+    parser.add_argument("--per-gpu-eval-batch-size", type=int, default=16)
+    parser.add_argument("--max-length", type=int, default=512, help="Maximum sequence length for the tokenizer.")
+    parser.add_argument("--total-epoch", type=int, default=10)
+    parser.add_argument("--warmup-steps", type=int, default=2000)
+    parser.add_argument("--lr", type=float, default=3e-5)
+    parser.add_argument("--min-lr", type=float, default=5e-6)
+    parser.add_argument("--weight-decay", type=float, default=1e-4)
+    parser.add_argument("--beta1", type=float, default=0.9)
+    parser.add_argument("--beta2", type=float, default=0.99)
+    parser.add_argument("--eps", type=float, default=1e-6)
+    parser.add_argument("--adv-p", type=float, default=0.5, help="Probability of sampling adversarial data.")
+    parser.add_argument("--num-workers-eval", type=int, default=8, help="Reserved for compatibility.")
+
+    parser.add_argument("--lora-r", type=int, default=128)
+    parser.add_argument("--lora-alpha", type=int, default=256)
+    parser.add_argument("--lora-dropout", type=float, default=0.0)
+
+    parser.add_argument("--freeze-layer", type=int, default=0, help="Number of initial encoder layers to freeze.")
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--adapter-path", type=Path, default=None, help="Optional path to resume LoRA training from.")
+    parser.add_argument("--pooling", type=str, default="max", choices=("max", "average", "cls"))
+
+    parser.add_argument("--lora", dest="lora", action="store_true", help="Enable LoRA adapters.")
+    parser.add_argument("--no-lora", dest="lora", action="store_false", help="Disable LoRA adapters.")
+    parser.set_defaults(lora=True)
+
+    parser.add_argument("--freeze-embedding-layer", dest="freeze_embedding_layer", action="store_true")
+    parser.add_argument("--no-freeze-embedding-layer", dest="freeze_embedding_layer", action="store_false")
+    parser.set_defaults(freeze_embedding_layer=True)
+
+    parser.add_argument("--adversarial", dest="adversarial", action="store_true")
+    parser.add_argument("--no-adversarial", dest="adversarial", action="store_false")
+    parser.set_defaults(adversarial=True)
+
+    parser.add_argument("--include-attack", dest="include_attack", action="store_true")
+    parser.add_argument("--no-include-attack", dest="include_attack", action="store_false")
+    parser.set_defaults(include_attack=True)
+
+    parser.add_argument("--has-mix", dest="has_mix", action="store_true")
+    parser.add_argument("--no-has-mix", dest="has_mix", action="store_false")
+    parser.set_defaults(has_mix=True)
+
+    parser.add_argument("--deepspeed", action="store_true", help="Use DeepSpeed strategy when multiple GPUs are available.")
+
+    return parser
+
+
+def train(args: argparse.Namespace) -> None:
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+    torch.set_float32_matmul_precision("medium")
+
+    if args.device_num > 1:
+        if args.deepspeed:
+            strategy = DeepSpeedStrategy()
+        else:
+            strategy = DDPStrategy(find_unused_parameters=True)
+        fabric = Fabric(accelerator="cuda", precision="bf16-mixed", devices=args.device_num, strategy=strategy)
+    else:
+        fabric = Fabric(accelerator="cuda", precision="bf16-mixed", devices=args.device_num)
+
+    fabric.launch()
+
+    experiment_paths = ExperimentPaths(root=Path(args.output_dir), runs=Path(args.runs_dir))
+    if fabric.global_rank == 0:
+        experiment_paths.root.mkdir(parents=True, exist_ok=True)
+        experiment_paths.runs.mkdir(parents=True, exist_ok=True)
+    fabric.barrier()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    collate_fn = _build_collate_fn(tokenizer, args.max_length)
+
+    model = SimCLR_Tree(args, fabric).train()
+
+    data_path = load_datapath(
+        str(args.path),
+        include_adversarial=args.adversarial,
+        dataset_name=args.dataset_name,
+        include_attack=args.include_attack,
+    )[args.dataset]
+
+    train_dataset = SCLDataset(
+        data_path,
+        fabric,
+        tokenizer,
+        name2id=model.names2id,
+        has_mix=args.has_mix,
+        adv_p=args.adv_p,
+    )
+
+    passages_dataloader = DataLoader(
+        train_dataset,
+        batch_size=args.per_gpu_batch_size,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        shuffle=True,
+        drop_last=True,
+        collate_fn=collate_fn,
+    )
+
+    model.train()
+    if args.freeze_embedding_layer:
+        for name, param in model.model.named_parameters():
+            if "emb" in name or "model.pooler" in name:
+                param.requires_grad = False
+            if args.freeze_layer > 0:
+                for i in range(args.freeze_layer):
+                    if f"encoder.layer.{i}." in name:
+                        param.requires_grad = False
+
+    model = torch.compile(model)
+    if fabric.global_rank == 0:
+        print("Model has been initialized!")
+        for name, param in model.model.named_parameters():
+            print(name, param.requires_grad)
+
+    passages_dataloader = fabric.setup_dataloaders(passages_dataloader, use_distributed_sampler=False)
+    if fabric.global_rank == 0:
+        print("DataLoader has been initialized!")
+
+    if fabric.global_rank == 0:
+        writer = SummaryWriter(str(experiment_paths.runs))
+        print(f"Save dir is {args.output_dir}")
+        opt_dict = vars(args)
+        opt_dict["output_dir"] = str(args.output_dir)
+        with open(Path(args.output_dir) / "config.yaml", "w", encoding="utf-8") as file:
+            yaml.dump(opt_dict, file, sort_keys=False)
+    else:
+        writer = None
+
+    experiment_dir = experiment_paths.root
+
+    num_batches_per_epoch = len(passages_dataloader)
+    warmup_steps = args.warmup_steps
+    lr = args.lr
+    total_steps = args.total_epoch * num_batches_per_epoch - warmup_steps
+
+    optimizer = optim.AdamW(
+        filter(lambda p: p.requires_grad, model.parameters()),
+        lr=args.lr,
+        betas=(args.beta1, args.beta2),
+        eps=args.eps,
+        weight_decay=args.weight_decay,
+    )
+
+    schedule = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, total_steps, eta_min=args.min_lr)
+    model, optimizer = fabric.setup(model, optimizer)
+
+    if fabric.global_rank == 0:
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                print(name, param.requires_grad)
+
+    for epoch in range(args.total_epoch):
+        model.train()
+        avg_loss = 0.0
+        iterator = enumerate(passages_dataloader)
+        if fabric.global_rank == 0:
+            iterator = tqdm(iterator, total=len(passages_dataloader))
+            print(("\n" + "%11s" * 5) % ("Epoch", "GPU_mem", "loss1", "Avgloss", "lr"))
+        for i, batch in iterator:
+            current_step = epoch * num_batches_per_epoch + i
+            if current_step < warmup_steps:
+                current_lr = lr * current_step / max(warmup_steps, 1)
+                for param_group in optimizer.param_groups:
+                    param_group["lr"] = current_lr
+            current_lr = optimizer.param_groups[0]["lr"]
+
+            encoded_batch, label, write_model = batch
+            loss, loss_classify = model(encoded_batch, write_model)
+
+            avg_loss = (avg_loss * i + loss.item()) / (i + 1)
+            fabric.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+            if current_step >= warmup_steps:
+                schedule.step()
+
+            mem = f"{torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0:.3g}G"
+            if fabric.global_rank == 0:
+                iterator.set_description(
+                    ("%11s" * 2 + "%11.4g" * 3)
+                    % (f"{epoch + 1}/{args.total_epoch}", mem, loss_classify.item(), avg_loss, current_lr)
+                )
+                if writer and current_step % 10 == 0:
+                    writer.add_scalar("lr", current_lr, current_step)
+                    writer.add_scalar("loss", loss.item(), current_step)
+                    writer.add_scalar("avg_loss", avg_loss, current_step)
+                    writer.add_scalar("loss_classify", loss_classify.item(), current_step)
+
+        if fabric.global_rank == 0:
+            checkpoint_dir = experiment_dir / f"epoch_{epoch:02d}"
+            model.save_pretrained(str(checkpoint_dir), save_tokenizer=(epoch == 0))
+            print(f"Saved adapter checkpoint to {checkpoint_dir}", flush=True)
+
+            last_dir = experiment_dir / "last"
+            model.save_pretrained(str(last_dir), save_tokenizer=False)
+            print(f"Updated latest checkpoint at {last_dir}", flush=True)
+
+        fabric.barrier()
+
+    if writer:
+        writer.flush()
+        writer.close()
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    parser = build_argument_parser()
+    args = parser.parse_args(argv)
+    experiment_paths = _prepare_output_dir(
+        args.output_dir, args.experiment_name, resume=args.resume, create_dirs=False
+    )
+    args.output_dir = str(experiment_paths.root)
+    args.runs_dir = str(experiment_paths.runs)
+    train(args)
+
+
+__all__ = ["build_argument_parser", "main", "train"]
+
+
+if __name__ == "__main__":
+    main()