admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-29 04:36:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(logging): standardize logging with getLogger(__name__) in scripts (#4192)

Browse Source 
* refactor(logging): replace print with logger in lerobot_info

* refactor(logging): replace print with logger in convert_dataset_v21_to_v30

* refactor(logging): replace print with logger in lerobot_annotate

* refactor(logging): replace print with logger in lerobot_dataset_viz

* refactor(logging): replace print with logger in lerobot_eval

* refactor(logging): replace print with logger in lerobot_find_cameras

* refactor(logging): replace print with logger in lerobot_find_joint_limits

* refactor(logging): replace print with logger in lerobot_find_port

* refactor(logging): replace print with logger in lerobot_imgtransform_viz

* refactor(logging): replace print with logger in lerobot_setup_can

* refactor(logging): replace print with logger in lerobot_teleoperate

* refactor(logging): replace print with logger in lerobot_train_tokenizer

* fix(logging): preserve CLI output semantics

---------

Co-authored-by: ailisilob <2248345706@qq.com>
This commit is contained in:
Steven Palma
2026-07-28 14:42:38 +02:00
committed by GitHub
parent a855570097
commit 4af7c70664
6 changed files with 82 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/convert_dataset_v21_to_v30.py
+6 -4
View File
@@ -94,6 +94,8 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
from lerobot.utils.constants import HF_LEROBOT_HOME from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.utils import flatten_dict, init_logging from lerobot.utils.utils import flatten_dict, init_logging
logger = logging.getLogger(__name__)
V21 = "v2.1" V21 = "v2.1"
V30 = "v3.0" V30 = "v3.0"
@@ -476,11 +478,11 @@ def convert_dataset(
# First check if the dataset already has a v3.0 version # First check if the dataset already has a v3.0 version
if root is None and not force_conversion: if root is None and not force_conversion:
try: try:
print("Trying to download v3.0 version of the dataset from the hub...") logger.info("Trying to download v3.0 version of the dataset from the hub...")
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id) snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
return return
except Exception: except Exception:
print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.") logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
# Set root based on whether local dataset path is provided # Set root based on whether local dataset path is provided
use_local_dataset = False use_local_dataset = False
@@ -488,7 +490,7 @@ def convert_dataset(
if root.exists(): if root.exists():
validate_local_dataset_version(root) validate_local_dataset_version(root)
use_local_dataset = True use_local_dataset = True
print(f"Using local dataset at {root}") logger.info(f"Using local dataset at {root}")
old_root = root.parent / f"{root.name}_old" old_root = root.parent / f"{root.name}_old"
new_root = root.parent / f"{root.name}_v30" new_root = root.parent / f"{root.name}_v30"
@@ -523,7 +525,7 @@ def convert_dataset(
try: try:
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset") hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
except (HTTPError, RevisionNotFoundError) as e: except (HTTPError, RevisionNotFoundError) as e:
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})") logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
pass pass
hub_api.delete_files( hub_api.delete_files(
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"], delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
src/lerobot/scripts/lerobot_annotate.py
+6 -7
View File
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
repo_id = cfg.new_repo_id or cfg.repo_id repo_id = cfg.new_repo_id or cfg.repo_id
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)" commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
api = HfApi() api = HfApi()
print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True) logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
api.create_repo( api.create_repo(
repo_id=repo_id, repo_id=repo_id,
repo_type="dataset", repo_type="dataset",
private=cfg.push_private, private=cfg.push_private,
exist_ok=True, exist_ok=True,
) )
print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True) logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
commit_info = api.upload_folder( commit_info = api.upload_folder(
folder_path=str(root), folder_path=str(root),
repo_id=repo_id, repo_id=repo_id,
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
# at the source dataset; a fresh card is generated below instead. # at the source dataset; a fresh card is generated below instead.
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"], ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
) )
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True) logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
dataset_info = load_info(root) dataset_info = load_info(root)
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id) card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
@@ -200,14 +200,13 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
with suppress(RevisionNotFoundError): with suppress(RevisionNotFoundError):
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset") api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
api.create_tag(**tag_kwargs) api.create_tag(**tag_kwargs)
print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True) logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
except Exception as exc: # noqa: BLE001 except Exception as exc: # noqa: BLE001
print( logger.warning(
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. " f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. " "Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
"Run: from huggingface_hub import HfApi; " "Run: from huggingface_hub import HfApi; "
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)", f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
flush=True,
) )
src/lerobot/scripts/lerobot_dataset_viz.py
+3 -1
View File
@@ -89,6 +89,8 @@ from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
DEFAULT_FOXGLOVE_PORT = 8765 DEFAULT_FOXGLOVE_PORT = 8765
DEFAULT_RERUN_PORT = 9090 DEFAULT_RERUN_PORT = 9090
@@ -299,7 +301,7 @@ def visualize_dataset(
while True: while True:
time.sleep(1) time.sleep(1)
except KeyboardInterrupt: except KeyboardInterrupt:
print("Ctrl-C received. Exiting.") logger.info("Ctrl-C received. Exiting.")
def main(): def main():
src/lerobot/scripts/lerobot_eval.py
+6 -4
View File
@@ -95,6 +95,8 @@ from lerobot.utils.utils import (
inside_slurm, inside_slurm,
) )
logger = logging.getLogger(__name__)
def _env_features_to_dataset_features(env_features: dict) -> dict: def _env_features_to_dataset_features(env_features: dict) -> dict:
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
@@ -796,13 +798,13 @@ def eval_main(cfg: EvalPipelineConfig):
recording_repo_id=cfg.eval.recording_repo_id, recording_repo_id=cfg.eval.recording_repo_id,
recording_private=cfg.eval.recording_private, recording_private=cfg.eval.recording_private,
) )
print("Overall Aggregated Metrics:") logger.info("Overall Aggregated Metrics:")
print(info["overall"]) logger.info(info["overall"])
# Print per-suite stats # Print per-suite stats
for task_group, task_group_info in info.items(): for task_group, task_group_info in info.items():
print(f"\nAggregated Metrics for {task_group}:") logger.info(f"\nAggregated Metrics for {task_group}:")
print(task_group_info) logger.info(task_group_info)
# Close all vec envs # Close all vec envs
close_envs(envs) close_envs(envs)
src/lerobot/scripts/lerobot_find_cameras.py
+3
View File
@@ -40,6 +40,7 @@ from PIL import Image
from lerobot.cameras import ColorMode from lerobot.cameras import ColorMode
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -285,6 +286,8 @@ def save_images_from_all_cameras(
def main(): def main():
init_logging()
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images." description="Unified camera utility script for listing cameras and capturing images."
) )
src/lerobot/scripts/lerobot_train_tokenizer.py
+58 -56
View File
@@ -45,6 +45,7 @@ lerobot-train-tokenizer \
""" """
import json import json
import logging
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@@ -63,6 +64,9 @@ else:
from lerobot.configs import NormalizationMode, parser from lerobot.configs import NormalizationMode, parser
from lerobot.datasets import LeRobotDataset from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@dataclass @dataclass
@@ -274,11 +278,8 @@ def process_episode(args):
return action_chunks return action_chunks
except Exception as e: except Exception:
print(f"Error processing episode {ep_idx}: {e}") logger.exception("Error processing episode %s", ep_idx)
import traceback
traceback.print_exc()
return None return None
@@ -300,10 +301,10 @@ def train_fast_tokenizer(
Returns: Returns:
Trained FAST tokenizer Trained FAST tokenizer
""" """
print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...") logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
print(f"Action chunk shape: {action_chunks.shape}") logger.info(f"Action chunk shape: {action_chunks.shape}")
print(f"Vocab size: {vocab_size}") logger.info(f"Vocab size: {vocab_size}")
print(f"DCT scale: {scale}") logger.info(f"DCT scale: {scale}")
# download the tokenizer source code (not pretrained weights) # download the tokenizer source code (not pretrained weights)
# we'll train a new tokenizer on our own data # we'll train a new tokenizer on our own data
@@ -314,7 +315,7 @@ def train_fast_tokenizer(
# train the new tokenizer on our action data using .fit() # train the new tokenizer on our action data using .fit()
# this trains the BPE tokenizer on DCT coefficients # this trains the BPE tokenizer on DCT coefficients
print("Training new tokenizer (this may take a few minutes)...") logger.info("Training new tokenizer (this may take a few minutes)...")
tokenizer = base_tokenizer.fit( tokenizer = base_tokenizer.fit(
action_data_list, action_data_list,
scale=scale, scale=scale,
@@ -322,21 +323,21 @@ def train_fast_tokenizer(
time_horizon=action_chunks.shape[1], # action_horizon time_horizon=action_chunks.shape[1], # action_horizon
action_dim=action_chunks.shape[2], # encoded dimensions action_dim=action_chunks.shape[2], # encoded dimensions
) )
print("✓ Tokenizer training complete!") logger.info("✓ Tokenizer training complete!")
# validate it works # validate it works
sample_chunk = action_chunks[0] sample_chunk = action_chunks[0]
encoded = tokenizer(sample_chunk[None])[0] encoded = tokenizer(sample_chunk[None])[0]
if isinstance(encoded, list): if isinstance(encoded, list):
encoded = np.array(encoded) encoded = np.array(encoded)
print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}") logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
return tokenizer return tokenizer
def compute_compression_stats(tokenizer, action_chunks: np.ndarray): def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"""Compute compression statistics.""" """Compute compression statistics."""
print("\nComputing compression statistics...") logger.info("\nComputing compression statistics...")
# sample for stats (use max 1000 chunks for speed) # sample for stats (use max 1000 chunks for speed)
sample_size = min(1000, len(action_chunks)) sample_size = min(1000, len(action_chunks))
@@ -366,12 +367,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"max_token_length": float(np.max(token_lengths)), "max_token_length": float(np.max(token_lengths)),
} }
print("Compression Statistics:") logger.info("Compression Statistics:")
print(f" Average compression ratio: {stats['compression_ratio']:.2f}x") logger.info(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
print(f" Mean token length: {stats['mean_token_length']:.1f}") logger.info(f" Mean token length: {stats['mean_token_length']:.1f}")
print(f" P99 token length: {stats['p99_token_length']:.0f}") logger.info(f" P99 token length: {stats['p99_token_length']:.0f}")
print(f" Min token length: {stats['min_token_length']:.0f}") logger.info(f" Min token length: {stats['min_token_length']:.0f}")
print(f" Max token length: {stats['max_token_length']:.0f}") logger.info(f" Max token length: {stats['max_token_length']:.0f}")
return stats return stats
@@ -385,9 +386,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
cfg: TokenizerTrainingConfig dataclass with all configuration parameters cfg: TokenizerTrainingConfig dataclass with all configuration parameters
""" """
# load dataset # load dataset
print(f"Loading dataset: {cfg.repo_id}") logger.info(f"Loading dataset: {cfg.repo_id}")
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root) dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames") logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
# parse normalization mode # parse normalization mode
try: try:
@@ -397,7 +398,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
f"Invalid normalization_mode: {cfg.normalization_mode}. " f"Invalid normalization_mode: {cfg.normalization_mode}. "
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}" f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
) from err ) from err
print(f"Normalization mode: {norm_mode.value}") logger.info(f"Normalization mode: {norm_mode.value}")
# parse encoded dimensions # parse encoded dimensions
encoded_dim_ranges = [] encoded_dim_ranges = []
@@ -406,38 +407,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_dim_ranges.append((start, end)) encoded_dim_ranges.append((start, end))
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges) total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}") logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
# parse relative dimensions # parse relative dimensions
relative_dim_list = None relative_dim_list = None
if cfg.relative_dims is not None and cfg.relative_dims.strip(): if cfg.relative_dims is not None and cfg.relative_dims.strip():
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")] relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
print(f"Relative dimensions: {relative_dim_list}") logger.info(f"Relative dimensions: {relative_dim_list}")
else: else:
print("No relative dimensions specified") logger.info("No relative dimensions specified")
print(f"Use relative transform: {cfg.use_relative_transform}") logger.info(f"Use relative transform: {cfg.use_relative_transform}")
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0): if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
print( logger.warning(
"Warning: use_relative_transform=True but no relative_dims specified. " "Warning: use_relative_transform=True but no relative_dims specified. "
"No relative transform will be applied." "No relative transform will be applied."
) )
print(f"Action horizon: {cfg.action_horizon}") logger.info(f"Action horizon: {cfg.action_horizon}")
print(f"State key: {cfg.state_key}") logger.info(f"State key: {cfg.state_key}")
# determine episodes to process # determine episodes to process
num_episodes = dataset.num_episodes num_episodes = dataset.num_episodes
if cfg.max_episodes is not None: if cfg.max_episodes is not None:
num_episodes = min(cfg.max_episodes, num_episodes) num_episodes = min(cfg.max_episodes, num_episodes)
print(f"Processing {num_episodes} episodes...") logger.info(f"Processing {num_episodes} episodes...")
# process episodes sequentially (to avoid pickling issues with dataset) # process episodes sequentially (to avoid pickling issues with dataset)
all_chunks = [] all_chunks = []
for ep_idx in range(num_episodes): for ep_idx in range(num_episodes):
if ep_idx % 10 == 0: if ep_idx % 10 == 0:
print(f" Processing episode {ep_idx}/{num_episodes}...") logger.info(f" Processing episode {ep_idx}/{num_episodes}...")
chunks = process_episode( chunks = process_episode(
( (
@@ -455,19 +456,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# concatenate all chunks # concatenate all chunks
all_chunks = np.concatenate(all_chunks, axis=0) all_chunks = np.concatenate(all_chunks, axis=0)
print(f"Collected {len(all_chunks)} action chunks") logger.info(f"Collected {len(all_chunks)} action chunks")
# extract only encoded dimensions FIRST (before normalization) # extract only encoded dimensions FIRST (before normalization)
encoded_chunks = [] encoded_chunks = []
for start, end in encoded_dim_ranges: for start, end in encoded_dim_ranges:
encoded_chunks.append(all_chunks[:, :, start:end]) encoded_chunks.append(all_chunks[:, :, start:end])
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded] encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions") logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
# apply normalization to encoded dimensions # apply normalization to encoded dimensions
print("\nBefore normalization - overall stats:") logger.info("\nBefore normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}") logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}") logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
# get normalization stats from dataset # get normalization stats from dataset
norm_stats = dataset.meta.stats norm_stats = dataset.meta.stats
@@ -489,9 +490,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_stats[stat_name] = stat_array[encoded_dim_indices] encoded_stats[stat_name] = stat_array[encoded_dim_indices]
if encoded_stats: if encoded_stats:
print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):") logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
for stat_name, stat_values in encoded_stats.items(): for stat_name, stat_values in encoded_stats.items():
print( logger.info(
f" {stat_name}: shape={stat_values.shape}, " f" {stat_name}: shape={stat_values.shape}, "
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]" f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
) )
@@ -499,27 +500,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# apply normalization based on mode # apply normalization based on mode
try: try:
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8) encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
print(f"\nApplied {norm_mode.value} normalization") logger.info(f"\nApplied {norm_mode.value} normalization")
except ValueError as e: except ValueError as e:
print(f"Warning: {e}. Using raw actions without normalization.") logger.warning(f"Warning: {e}. Using raw actions without normalization.")
print("\nAfter normalization - overall stats:") logger.info("\nAfter normalization - overall stats:")
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}") logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}") logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
print("\nPer-dimension stats (after normalization):") logger.info("\nPer-dimension stats (after normalization):")
for d in range(encoded_chunks.shape[-1]): for d in range(encoded_chunks.shape[-1]):
dim_data = encoded_chunks[:, :, d] dim_data = encoded_chunks[:, :, d]
print( logger.info(
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, " f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}" f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
) )
else: else:
print("Warning: Could not extract stats for encoded dimensions, using raw actions") logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
else: else:
print("Warning: No normalization stats found in dataset, using raw actions") logger.warning("Warning: No normalization stats found in dataset, using raw actions")
print(f"Encoded chunks shape: {encoded_chunks.shape}") logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
# train FAST tokenizer # train FAST tokenizer
tokenizer = train_fast_tokenizer( tokenizer = train_fast_tokenizer(
@@ -561,8 +562,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
with open(output_path / "metadata.json", "w") as f: with open(output_path / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2) json.dump(metadata, f, indent=2)
print(f"\nSaved FAST tokenizer to {output_path}") logger.info(f"\nSaved FAST tokenizer to {output_path}")
print(f"Metadata: {json.dumps(metadata, indent=2)}") logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
# push to Hugging Face Hub if requested # push to Hugging Face Hub if requested
if cfg.push_to_hub: if cfg.push_to_hub:
@@ -570,10 +571,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
hub_repo_id = cfg.hub_repo_id hub_repo_id = cfg.hub_repo_id
if hub_repo_id is None: if hub_repo_id is None:
hub_repo_id = output_path.name hub_repo_id = output_path.name
print(f"\nNo hub_repo_id provided, using: {hub_repo_id}") logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}") logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
print(f" Private: {cfg.hub_private}") logger.info(f" Private: {cfg.hub_private}")
try: try:
# use the tokenizer's push_to_hub method # use the tokenizer's push_to_hub method
@@ -593,14 +594,15 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
commit_message="Upload tokenizer metadata", commit_message="Upload tokenizer metadata",
) )
print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}") logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
except Exception as e: except Exception as e:
print(f"Error pushing to hub: {e}") logger.error(f"Error pushing to hub: {e}")
print(" Make sure you're logged in with `huggingface-cli login`") logger.error(" Make sure you're logged in with `huggingface-cli login`")
def main(): def main():
"""CLI entry point that parses arguments and runs the tokenizer training.""" """CLI entry point that parses arguments and runs the tokenizer training."""
init_logging()
train_tokenizer() train_tokenizer()