Reduce chat data exposure in generation and dataset handling

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,

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))

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)

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,15 +193,23 @@ 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:
             response = client.chat.completions.create(
                 **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():