Hugging Face's logo

doanh25032004
/
video_gen_physics_backup

Model card Files
xet
 Community
video_gen_physics_backup
File size: 4,931 Bytes
ec0a9aa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
"""
Join task_labels.csv back to meta/episodes.jsonl, filter to the first
NUM_CHUNKS chunks, and sample SAMPLE_PER_LABEL episodes per label.

Usage:
    python scripts/join_labels_to_episodes.py

    # Optional flags:
    python scripts/join_labels_to_episodes.py \
        --episodes-jsonl   datasets/droid_1.0.1_20chunks/meta/episodes.jsonl \
        --task-labels-csv  scripts/task_labels.csv \
        --output-csv       scripts/episode_markovian_split.csv \
        --num-chunks       20 \
        --sample-per-label 130 \
        --seed             42
"""

import argparse
import json
from pathlib import Path

import pandas as pd


CHUNKS_SIZE = 1000  # episodes per chunk (from info.json)


def load_episodes(episodes_jsonl: Path, max_episode: int) -> pd.DataFrame:
    rows = []
    with open(episodes_jsonl) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            ep_idx = obj.get("episode_index", -1)
            if ep_idx < 0 or ep_idx >= max_episode:
                continue
            length = obj.get("length")
            tasks = obj.get("tasks", [])
            task_text = next((t for t in tasks if t and t.strip()), "")
            rows.append(
                {
                    "episode_index": ep_idx,
                    "length": length,
                    "task": task_text,
                }
            )
    return pd.DataFrame(rows)


def main() -> None:
    parser = argparse.ArgumentParser(description="Join task labels to episodes and sample")
    parser.add_argument(
        "--episodes-jsonl",
        type=Path,
        default=Path("datasets/droid_1.0.1_20chunks/meta/episodes.jsonl"),
    )
    parser.add_argument(
        "--task-labels-csv",
        type=Path,
        default=Path("scripts/task_labels.csv"),
    )
    parser.add_argument(
        "--output-csv",
        type=Path,
        default=Path("scripts/episode_markovian_split.csv"),
    )
    parser.add_argument(
        "--num-chunks",
        type=int,
        default=20,
        help="Only include episodes from the first N chunks",
    )
    parser.add_argument(
        "--sample-per-label",
        type=int,
        default=130,
        help="Number of episodes to sample per label (0 = keep all)",
    )
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    workspace = Path(__file__).parent.parent

    def resolve(p: Path) -> Path:
        return workspace / p if not p.is_absolute() else p

    episodes_jsonl = resolve(args.episodes_jsonl)
    task_labels_csv = resolve(args.task_labels_csv)
    output_csv = resolve(args.output_csv)
    output_csv.parent.mkdir(parents=True, exist_ok=True)

    max_episode = args.num_chunks * CHUNKS_SIZE
    print(f"[load] episodes (chunks 0–{args.num_chunks - 1}, episode_index < {max_episode})")
    episodes_df = load_episodes(episodes_jsonl, max_episode)
    print(f"       {len(episodes_df)} episodes loaded")

    print(f"[load] task labels: {task_labels_csv}")
    labels_df = pd.read_csv(task_labels_csv)
    print(f"       {len(labels_df)} task labels loaded")
    print(f"       label counts:\n{labels_df['label'].value_counts().to_string()}")

    # Join on task text (episodes.jsonl has no task_index field)
    merged = episodes_df.merge(
        labels_df[["task_index", "task", "label", "reason"]],
        on="task",
        how="left",
    )

    unmatched = merged["label"].isna().sum()
    if unmatched:
        print(f"[warn] {unmatched} episodes have no matching label (empty/unknown task)")
    merged["label"] = merged["label"].fillna("unknown")
    merged["reason"] = merged["reason"].fillna("")

    # Save full joined table before sampling
    full_csv = output_csv.with_name(output_csv.stem + "_full.csv")
    merged.sort_values("episode_index").to_csv(full_csv, index=False)
    print(f"\n[saved] full join → {full_csv}  ({len(merged)} rows)")
    print(merged["label"].value_counts().to_string())

    # Sample per label
    if args.sample_per_label > 0:
        sampled_parts = []
        for label, group in merged.groupby("label"):
            if label in ("unknown", "parse_error", "api_error"):
                continue
            n = min(args.sample_per_label, len(group))
            sampled_parts.append(group.sample(n=n, random_state=args.seed))
            print(f"[sample] {label}: {n}/{len(group)} episodes selected")

        sampled = pd.concat(sampled_parts).sort_values("episode_index").reset_index(drop=True)
        sampled.to_csv(output_csv, index=False)
        print(f"\n[saved] sampled → {output_csv}  ({len(sampled)} rows)")
        print(sampled["label"].value_counts().to_string())
    else:
        merged.sort_values("episode_index").to_csv(output_csv, index=False)
        print(f"\n[saved] {output_csv}  ({len(merged)} rows)")


if __name__ == "__main__":
    main()