Upload generate_validator_data.py with huggingface_hub

generate_validator_data.py

@@ -0,0 +1,335 @@
+#!/usr/bin/env python3
+"""Generate validator-matching AI + human text samples with resume support.
+
+Usage:
+  python3 generate_validator_data.py --output data.jsonl --n-samples 1000
+  python3 generate_validator_data.py --output data.jsonl --restart   # force fresh
+  python3 generate_validator_data.py --n-samples 1000 --dry-run      # show plan only
+"""
+import sys, os, json, time, logging, argparse, random
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logger = logging.getLogger(__name__)
+
+from validator_data_gen.config import MODELS, AI_IN_MIDDLE_PROB, N_SAMPLES, N_HUMAN_SAMPLES, N_AI_SAMPLES
+
+KEY_FILES = {
+    'groq': '/root/groqkey1',
+    'groq2': '/root/groqkey2',
+    'nvidia_nim': '/root/nvidianmikey',
+    'llm7': '/root/llm7key',
+    'sambanova': '/root/sambanovakey',
+    'deepseek': '/root/deepseekkey',
+    'gemini': '/root/geminikey',
+}
+TOKENIZER_NAME = 'pangram/editlens_roberta-large'
+
+
+def load_keys():
+    keys = {}
+    for name, path in KEY_FILES.items():
+        if os.path.exists(path):
+            keys[name] = open(path).read().strip()
+    return keys
+
+
+def compute_targets(n_ai, n_human):
+    """Distribute n_ai samples across model slots matching validator ratios."""
+    unique = [m for m in MODELS if not m.get('_dup')]
+    mid_models = [m for m in unique if m.get('in_the_middle')]
+    n_mid = int(n_ai * AI_IN_MIDDLE_PROB)
+
+    targets = []
+    if mid_models:
+        per_mid = n_mid // len(mid_models)
+        extra_mid = n_mid - per_mid * len(mid_models)
+        for i, m in enumerate(mid_models):
+            cnt = per_mid + (1 if i < extra_mid else 0)
+            if cnt > 0:
+                targets.append({'name': m['name'], 'type': 'ai_in_middle', 'target': cnt, 'text_mode': m.get('text_mode', False)})
+
+    n_full = n_ai - n_mid
+    if unique:
+        per_full = n_full // len(unique)
+        extra_full = n_full - per_full * len(unique)
+        for i, m in enumerate(unique):
+            cnt = per_full + (1 if i < extra_full else 0)
+            if cnt > 0:
+                targets.append({'name': m['name'], 'type': 'ai_full', 'target': cnt, 'text_mode': m.get('text_mode', False)})
+
+    return targets, n_human
+
+
+def recount_output(path):
+    """Count complete JSON lines in output file, return count and fix truncated last line."""
+    if not os.path.exists(path):
+        return 0
+    count = 0
+    with open(path, 'r') as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                json.loads(line)
+                count += 1
+            except json.JSONDecodeError:
+                fix_path = path + '.fix'
+                with open(path, 'r') as r, open(fix_path, 'w') as w:
+                    for i, l in enumerate(r):
+                        if i < count:
+                            w.write(l)
+                os.replace(fix_path, path)
+                logger.warning(f'Truncated malformed last line, output now has {count} records')
+                break
+    return count
+
+
+def load_state(path):
+    if not os.path.exists(path):
+        return None
+    with open(path) as f:
+        return json.load(f)
+
+
+def write_state(state, path):
+    tmp = path + '.tmp'
+    with open(tmp, 'w') as f:
+        json.dump(state, f, indent=2)
+        f.flush()
+        os.fsync(f.fileno())
+    os.replace(tmp, path)
+
+
+def _load_human_texts(n=600):
+    """Download human text samples from HF dataset."""
+    from datasets import load_dataset
+    ds = load_dataset('cc_news', split='train', streaming=True).take(n + 100)
+    texts = []
+    for sample in ds:
+        text = sample.get('text') or sample.get('title', '') + '\n' + sample.get('description', '')
+        text = text.strip()
+        if len(text.split()) >= 100:
+            texts.append(text)
+        if len(texts) >= n:
+            break
+    if len(texts) < n:
+        logger.warning(f'Only got {len(texts)} human texts (wanted {n})')
+    logger.info(f'Loaded {len(texts)} human texts from cc_news')
+    return texts
+
+
+def text_source(texts, min_len=500):
+    idx = list(range(len(texts)))
+    random.shuffle(idx)
+    i = 0
+    while True:
+        text = texts[idx[i % len(texts)]]
+        if len(text) < min_len:
+            i += 1
+            continue
+        yield text
+        i += 1
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output', default='generated_data.jsonl')
+    parser.add_argument('--n-samples', type=int, default=N_SAMPLES)
+    parser.add_argument('--restart', action='store_true')
+    parser.add_argument('--dry-run', action='store_true')
+    args = parser.parse_args()
+
+    keys = load_keys()
+    if not keys:
+        logger.error('No API keys found')
+        sys.exit(1)
+
+    from validator_data_gen.api_hub import APIHub
+    from validator_data_gen.model_map import build_providers
+    from validator_data_gen.replicate import generate_one_sample
+
+    providers = build_providers(keys)
+    hub = APIHub(providers)
+    logger.info(f'Providers: {[p.name for p in providers]}')
+
+    # Load human text source from HF datasets (no local cache needed)
+    logger.info('Loading human text source from HF datasets...')
+    human_texts = _load_human_texts(n=2000)
+    logger.info(f'Loaded {len(human_texts)} human texts')
+
+    # Tokenizer not needed for generation — prep_training_data.py handles that
+
+    n_ai = args.n_samples
+    n_human_target = max(1, int(n_ai * N_HUMAN_SAMPLES / N_AI_SAMPLES))
+    targets, _ = compute_targets(n_ai, n_human_target)
+
+    if args.dry_run:
+        logger.info(f'=== DRY RUN: {n_ai} AI + {n_human_target} human ===')
+        by_type = {}
+        for t in targets:
+            by_type.setdefault(t['type'], []).append(t)
+        for ttype, items in by_type.items():
+            logger.info(f'  {ttype}: {sum(i["target"] for i in items)} samples across {len(items)} models')
+            for i in items:
+                logger.info(f'    {i["name"]}: {i["target"]}')
+        logger.info(f'  human: {n_human_target}')
+        return
+
+    state_path = args.output.rsplit('.', 1)[0] + '_state.json'
+    output_fd = None
+
+    # Resume or fresh start
+    if args.restart or not os.path.exists(state_path):
+        existing = recount_output(args.output)
+        if existing > 0 and not args.restart:
+            logger.info(f'Found {existing} existing records, resuming')
+            # will reconcile below
+            state = load_state(state_path) or {}
+        else:
+            state = {
+                'output_path': args.output, 'n_ai': n_ai, 'n_human_target': n_human_target,
+                'targets': targets, 'human_done': 0, 'total_done': 0, 'started_at': time.time(),
+                'version': 2,
+            }
+            # start fresh
+            if args.restart and existing > 0:
+                logger.info(f'Restart forced, discarding {existing} existing records')
+                os.remove(args.output)
+    else:
+        state = load_state(state_path)
+        if state is None:
+            logger.error(f'State file {state_path} corrupted')
+            sys.exit(1)
+        existing = recount_output(args.output)
+        logger.info(f'Resuming: state says {state["total_done"]} done, output has {existing} records')
+
+    # Reconcile: state may claim more done than output actually has (crash after output write but before state update)
+    if state['total_done'] > existing:
+        diff = state['total_done'] - existing
+        logger.warning(f'State ahead by {diff} — crash occurred after output write. Correcting state.')
+        state['total_done'] = existing
+        # rebuild per-target done from output file
+        if existing > 0:
+            done_map = {}
+            with open(args.output) as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        r = json.loads(line)
+                        if r['type'] != 'human':
+                            key = (r['model'], r['type'])
+                            done_map[key] = done_map.get(key, 0) + 1
+                        else:
+                            state['human_done'] = state.get('human_done', 0) + 1
+                    except json.JSONDecodeError:
+                        continue
+            for t in state['targets']:
+                t['done'] = done_map.get((t['name'], t['type']), 0)
+
+    # Open output for appending
+    output_fd = open(args.output, 'a')
+    os.fsync(output_fd.fileno())
+
+    # Human text source
+    human_gen = text_source(human_texts, min_len=300)
+
+    # Primary generation loop
+    gen_start = time.time()
+    consecutive_model_failures = 0
+
+    for ti, target in enumerate(state['targets']):
+        model_name = target['name']
+        gen_type = target['type']
+        target_cnt = target['target']
+        done = target.get('done', 0)
+        remaining = target_cnt - done
+
+        if remaining <= 0:
+            continue
+
+        logger.info(f'[{ti+1}/{len(state["targets"])}] {model_name} ({gen_type}): {done}/{target_cnt} done, {remaining} remaining')
+        model_fails = 0
+
+        for j in range(remaining):
+            sample = None
+            try:
+                src_text = next(human_gen)
+                prompt = src_text[:int(len(src_text) * random.uniform(0.25, 0.75))]
+                sample = generate_one_sample(hub, model_name, gen_type, target.get('text_mode', False), src_text, prompt)
+            except Exception as e:
+                logger.warning(f'{model_name}/{gen_type} attempt {j}: {e}')
+                model_fails += 1
+                consecutive_model_failures += 1
+                if model_fails >= 5:
+                    logger.warning(f'{model_name} failed {model_fails} times consecutively, skipping')
+                    break
+                if consecutive_model_failures >= 10:
+                    logger.warning('Too many consecutive failures across models, exiting')
+                    break
+                continue
+
+            if sample is None:
+                model_fails += 1
+                consecutive_model_failures += 1
+                if model_fails >= 5:
+                    break
+                j -= 1  # retry same index
+                continue
+
+            consecutive_model_failures = 0
+            model_fails = 0
+
+            # Atomic write: output → fsync → state update → state write
+            output_fd.write(json.dumps(sample) + '\n')
+            output_fd.flush()
+            os.fsync(output_fd.fileno())
+
+            # Confirm it was written (read-back check)
+            state['total_done'] += 1
+            target['done'] = target.get('done', 0) + 1
+            write_state(state, state_path)
+
+            if (j + 1) % 10 == 0:
+                elapsed = time.time() - gen_start
+                rate = state['total_done'] / max(elapsed, 1)
+                logger.info(f'  {model_name}: {target["done"]}/{target_cnt} done, total={state["total_done"]}, rate={rate:.2f}/s')
+
+        target['finalized'] = True
+        write_state(state, state_path)
+
+        if consecutive_model_failures >= 10:
+            logger.warning('Too many consecutive failures, stopping generation')
+            break
+
+    # Human phase (no API calls, just from cache)
+    human_remaining = n_human_target - state.get('human_done', 0)
+    if human_remaining > 0:
+        logger.info(f'Generating {human_remaining} human samples from cache...')
+        for j in range(human_remaining):
+            src_text = next(human_gen)
+            labels = [0] * len(src_text.split())
+            sample = {'text': src_text, 'text_raw': src_text, 'labels': labels,
+                      'model': 'human', 'params': {}, 'type': 'human', 'augmentations': []}
+            output_fd.write(json.dumps(sample) + '\n')
+            output_fd.flush()
+            os.fsync(output_fd.fileno())
+            state['human_done'] = state.get('human_done', 0) + 1
+            state['total_done'] += 1
+            write_state(state, state_path)
+
+    output_fd.close()
+    state['status'] = 'completed'
+    state['elapsed'] = time.time() - gen_start
+    write_state(state, state_path)
+
+    logger.info(f'Done! {state["total_done"]} samples in {state["elapsed"]:.0f}s')
+    n_ai_done = sum(t.get('done', 0) for t in state['targets'])
+    logger.info(f'  AI: {n_ai_done}, Human: {state.get("human_done", 0)}')
+
+
+if __name__ == '__main__':
+    main()