diff --git a/deepspec/data/jsonl_dataset.py b/deepspec/data/jsonl_dataset.py index a3c22191..bb5e1729 100644 --- a/deepspec/data/jsonl_dataset.py +++ b/deepspec/data/jsonl_dataset.py @@ -4,19 +4,35 @@ import os import pickle from bisect import bisect_right +from pathlib import Path import torch from tqdm import tqdm CACHE_DIR = os.path.expanduser("~/.cache/deepspec") +_CACHE_DIR_ENV = "DEEPSPEC_CACHE_DIR" +_DISABLE_INDEX_CACHE_ENV = "DEEPSPEC_DISABLE_JSONL_INDEX_CACHE" + + +def _is_disabled_env(value): + return str(value).strip().lower() in {"1", "true", "yes", "on"} class JsonLineDataset(torch.utils.data.Dataset): - def __init__(self, data_paths): + def __init__(self, data_paths, *, cache_dir=None, cache_index=True): super().__init__() self.data_paths = sorted(data_paths) - self.cache_dir = os.path.abspath(CACHE_DIR) - os.makedirs(self.cache_dir, exist_ok=True) + self.cache_index = bool(cache_index) and not _is_disabled_env( + os.environ.get(_DISABLE_INDEX_CACHE_ENV, "") + ) + resolved_cache_dir = ( + cache_dir + if cache_dir is not None + else os.environ.get(_CACHE_DIR_ENV, CACHE_DIR) + ) + self.cache_dir = os.path.abspath(os.path.expanduser(resolved_cache_dir)) + if self.cache_index: + os.makedirs(self.cache_dir, exist_ok=True) self.num_data_per_file = [] self.cum_counts = [0] self.files = [None] * len(self.data_paths) @@ -73,18 +89,27 @@ def _file_key(self, path): abspath = os.path.abspath(path) st = os.stat(abspath) mtime_ns = getattr(st, "st_mtime_ns", int(st.st_mtime * 1e9)) - return f"{os.path.normcase(abspath)}|{mtime_ns}" + path_hash = hashlib.blake2b( + os.path.normcase(abspath).encode("utf-8"), + digest_size=16, + ).hexdigest() + return f"pathhash={path_hash}|size={st.st_size}|mtime_ns={mtime_ns}" def _cache_path_from_key(self, file_key): + if not self.cache_index: + return None key_hash = hashlib.blake2b(file_key.encode("utf-8"), digest_size=16).hexdigest() return os.path.join(self.cache_dir, f"jsonlindex-{key_hash}.pkl") def _atomic_pickle_dump(self, obj, dst_path): + if dst_path is None: + return tmp_path = f"{dst_path}.tmp-{os.getpid()}" with open(tmp_path, "wb") as handle: pickle.dump(obj, handle, protocol=pickle.HIGHEST_PROTOCOL) handle.flush() os.fsync(handle.fileno()) + os.chmod(tmp_path, 0o600) os.replace(tmp_path, dst_path) def _build_all_line_starts(self): @@ -125,7 +150,7 @@ def _build_all_line_starts(self): self._atomic_pickle_dump( { "file_key": file_key, - "file_path": os.path.abspath(path), + "source_name": Path(path).name, "line_starts": starts, }, cache_path, diff --git a/deepspec/eval/base_evaluator.py b/deepspec/eval/base_evaluator.py index d094ec38..6f1ddf9d 100644 --- a/deepspec/eval/base_evaluator.py +++ b/deepspec/eval/base_evaluator.py @@ -26,13 +26,20 @@ DEFAULT_DATASET_ROOT = "./eval_datasets" +_DATASET_ROOT_ENV = "DEEPSPEC_EVAL_DATASET_ROOT" + + +def _resolve_dataset_root(dataset_root: str | None = None) -> str: + if dataset_root is not None: + return dataset_root + return os.environ.get(_DATASET_ROOT_ENV, DEFAULT_DATASET_ROOT) def load_and_process_dataset( data_name: str, - dataset_root: str = DEFAULT_DATASET_ROOT, + dataset_root: str | None = None, ): - dataset_path = Path(dataset_root) / f"{data_name}.jsonl" + dataset_path = Path(_resolve_dataset_root(dataset_root)) / f"{data_name}.jsonl" assert dataset_path.exists() rows = [] @@ -517,7 +524,10 @@ def run_dataset( max_samples: int | None, ) -> list[SimpleNamespace]: seed_all(int(self.args.seed)) - dataset = load_and_process_dataset(dataset_name) + dataset = load_and_process_dataset( + dataset_name, + dataset_root=getattr(self.args, "dataset_root", None), + ) if max_samples is not None and len(dataset) > max_samples: rng = random.Random(int(self.args.seed)) diff --git a/eval.py b/eval.py index a35e7ea3..f45781ef 100644 --- a/eval.py +++ b/eval.py @@ -1,6 +1,7 @@ from __future__ import annotations import argparse import json +import os import torch from transformers import AutoConfig from deepspec.eval.dspark import Gemma4DSparkEvaluator, Qwen3DSparkEvaluator @@ -41,6 +42,12 @@ def parse_args(): ) parser.add_argument("--tensorboard-dir", type=str, default=None) parser.add_argument("--step", type=int, default=None,help=("step for tensorboard logging"),) + parser.add_argument( + "--dataset-root", + type=str, + default="./eval_datasets", + help="Directory containing evaluation JSONL files.", + ) parser.add_argument("--seed", type=int, default=980406) args = parser.parse_args() args.tasks = list(TASKS) diff --git a/scripts/data/generate_train_data.py b/scripts/data/generate_train_data.py index 2ab156d0..7cd3e7dd 100644 --- a/scripts/data/generate_train_data.py +++ b/scripts/data/generate_train_data.py @@ -26,6 +26,15 @@ def parse_args(): parser.add_argument("--max-tokens", type=int, default=4096) parser.add_argument("--num-samples", type=int, default=None) parser.add_argument("--resume", action="store_true") + parser.add_argument("--privacy-test-flag", action="store_true") + parser.add_argument( + "--include-error-conversations", + action="store_true", + help=( + "Write full failed samples to the error JSONL. By default, message " + "content is redacted from errors to avoid retaining private chats." + ), + ) parser.add_argument("--is-reasoning-model", action="store_true") parser.add_argument("--is-gpt-oss", action="store_true") @@ -97,18 +106,81 @@ def build_query_kwargs(args, messages, max_tokens=None): return query_kwargs -def error_sample(sample, message): - sample["status"] = "error" - sample["error"] = message - return sample +SENSITIVE_SAMPLE_FIELDS = {"conversations", "messages", "turns"} + + +def _collect_sensitive_strings(value): + if isinstance(value, str): + return [value] if value else [] + if isinstance(value, dict): + sensitive = [] + for item in value.values(): + sensitive.extend(_collect_sensitive_strings(item)) + return sensitive + if isinstance(value, list): + sensitive = [] + for item in value: + sensitive.extend(_collect_sensitive_strings(item)) + return sensitive + return [] + + +def _redact_error_message(message, sample): + redacted = str(message) + for field in SENSITIVE_SAMPLE_FIELDS: + for sensitive_text in _collect_sensitive_strings(sample.get(field)): + if sensitive_text in redacted: + redacted = redacted.replace(sensitive_text, "[redacted]") + return redacted + + +def _redact_error_sample(sample): + redacted = {} + for key, value in sample.items(): + if key in SENSITIVE_SAMPLE_FIELDS: + continue + if isinstance(value, (int, float, bool)) or value is None: + redacted[key] = value + elif isinstance(value, str): + redacted[f"{key}_redacted"] = True + elif isinstance(value, list): + redacted[f"{key}_count"] = len(value) + redacted[f"{key}_redacted"] = True + elif isinstance(value, dict): + redacted[f"{key}_keys"] = sorted(str(item_key) for item_key in value.keys()) + redacted[f"{key}_redacted"] = True + else: + redacted[f"{key}_redacted"] = True + for field in SENSITIVE_SAMPLE_FIELDS: + if field not in sample: + continue + value = sample[field] + redacted[f"{field}_redacted"] = True + redacted[f"{field}_count"] = len(value) if isinstance(value, list) else None + return redacted + + +def error_sample(sample, message, *, include_conversations=False): + output_sample = sample if include_conversations else _redact_error_sample(sample) + output_sample["status"] = "error" + output_sample["error"] = _redact_error_message(message, sample) + return output_sample def call_sglang(args, server_address, sample, max_tokens=None): conversations = sample.get("conversations") if not conversations: - return error_sample(sample, "Missing conversations") + return error_sample( + sample, + "Missing conversations", + include_conversations=args.include_error_conversations, + ) if conversations[0].get("role") == "assistant": - return error_sample(sample, "Data starts with an assistant message") + return error_sample( + sample, + "Data starts with an assistant message", + include_conversations=args.include_error_conversations, + ) client = OpenAI(base_url=f"http://{server_address}/v1", api_key="None") regenerated = [] @@ -121,7 +193,11 @@ def call_sglang(args, server_address, sample, max_tokens=None): if role == "assistant": continue if role != "user": - return error_sample(sample, f"Invalid message role: {role}") + return error_sample( + sample, + f"Invalid message role: {role}", + include_conversations=args.include_error_conversations, + ) regenerated.append(message) try: @@ -129,7 +205,11 @@ def call_sglang(args, server_address, sample, max_tokens=None): **build_query_kwargs(args, regenerated, max_tokens=max_tokens) ) except Exception as exc: - return error_sample(sample, str(exc)) + return error_sample( + sample, + str(exc), + include_conversations=args.include_error_conversations, + ) response_message = { "role": "assistant", @@ -190,7 +270,7 @@ def validate_servers(args): try: _, result, elapsed = future.result() except Exception as exc: - result = {"status": "error", "error": str(exc)} + result = {"status": "error", "error": "server validation failed"} elapsed = 0.0 if result.get("status") == "success": @@ -239,6 +319,7 @@ def write_finished_result( sample = future.result() if sample["status"] == "error": error_handle.write(json.dumps(sample, ensure_ascii=False) + "\n") + error_handle.flush() stats["errors"] += 1 return @@ -267,6 +348,9 @@ def print_config(args): print(f" top_k: {args.top_k}") print(f" min_p: {args.min_p}") print(f" resume: {args.resume}") + print(f" include_error_conversations: {args.include_error_conversations}") + print(f" include_error_conversations: {args.include_error_conversations}") + print(f" include_error_conversations: {args.include_error_conversations}") def main():