Unfinished code

[tastelikefeet]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Write the detail information belongs to this PR.

Experiment results

Paste your experiment result here(if needed).

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive suite of scripts and modules for building and evaluating a RAG index of thinking traces, training embedding models, and running GRPO reinforcement learning on math datasets. It also includes framework-level updates to the twinkle library, such as Ray timeout support and double-sided clipping in GRPO loss. The review feedback highlights several critical issues, including potential TypeErrors due to unhandled None returns from template.encode, thread-safety concerns during concurrent LanceDB writes, silent data loss in Cosmopedia parsing, a logical bug overriding valid 0 timeouts in Ray collection, and missing resource cleanup in background thread pools.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Since some ablation runs are commented out in this script, the glob patterns (e.g., ablation_*_24k.jsonl) might not match any files. In bash, unmatched globs are left unexpanded as literal strings. Since set -euo pipefail is active, trying to read from a non-existent literal glob path will cause wc to fail and the script to exit prematurely. We should check if the file exists before processing it.

Suggested change

	for f in $OUTDIR/ablation_*_65k.jsonl $OUTDIR/ablation_*_24k.jsonl; do
	n=$(wc -l < "$f")
	for f in $OUTDIR/ablation_*_65k.jsonl $OUTDIR/ablation_*_24k.jsonl; do
	[ -f "$f" ] || continue
	n=$(wc -l < "$f")

[gemini-code-assist]

Since _process_batch is submitted to a ThreadPoolExecutor with multiple workers, _embed_and_insert will be executed concurrently by multiple threads. LanceDB's Table.add is not thread-safe and concurrent writes can lead to database corruption or write conflicts. Additionally, mutating indexed (a Python set) and nonlocal_counters concurrently from multiple threads is not thread-safe. We should serialize the embedding, database insertion, and counter updates using a threading.Lock.

    write_lock = threading.Lock()

    def _embed_and_insert(kept_rows: List[Dict[str, Any]]) -> None:
        """Phase 2+3: embed compressed texts and insert into LanceDB."""
        if not kept_rows:
            return
        with write_lock:
            anchor_emb = get_embeddings(
                emb_model, emb_template, [r['query_compressed'] for r in kept_rows], role='anchor')
            positive_emb = get_embeddings(
                emb_model, emb_template, [r['cot_compressed'] for r in kept_rows], role='positive')
            sims = (anchor_emb * positive_emb).sum(axis=1).astype(np.float32)
            to_insert: List[Dict[str, Any]] = []
            for idx, (r, sim_val) in enumerate(zip(kept_rows, sims)):
                tag = 'KEEP' if sim_val >= SIM_THRESHOLD else 'DROP'
                print(f'[{tag} sim={sim_val:.4f}] {r["source"][:24]} '
                      f'q={_short(r["query_raw"], 60)!r} '
                      f'cot={_short(r["cot_raw"], 60)!r}', flush=True)
                if sim_val < SIM_THRESHOLD:
                    nonlocal_counters['n_dropped_sim'] += 1
                    _log_miss(misses_path, misses_lock, {
                        'id': r['id'], 'source': r['source'], 'reason': 'sim_low',
                        'sim': float(sim_val),
                        'query_raw': r['query_raw'],
                        'cot_raw': r['cot_raw'],
                        'query_compressed': r['query_compressed'],
                        'cot_compressed': r['cot_compressed'],
                    })
                    continue
                to_insert.append({
                    'id': r['id'],
                    'vector': positive_emb[idx].tolist(),
                    'thinking_raw': r['cot_raw'],
                    'query_raw': r['query_raw'],
                    'cot_compressed': r['cot_compressed'],
                    'query_compressed': r['query_compressed'],
                    'source': r['source'],
                    'domain': DOMAIN_MAP.get(r['source'], 'mixed'),
                    'language': _detect_lang(r['cot_raw']),
                    'sim': float(sim_val),
                })
            if to_insert:
                tbl.add(to_insert)
                nonlocal_counters['n_kept'] += len(to_insert)
                indexed.update(r['id'] for r in to_insert)

[gemini-code-assist]

Since template.encode can now return None (as updated in src/twinkle/template/base.py), executing feat['labels'] = ... without a None check will raise a TypeError. We must guard against None before mutating feat and only append valid features.

Suggested change

	if role == 'anchor':
	feat = template.encode({'messages': _wrap_anchor(text)})
	feat['labels'] = [1]
	else:
	feat = template.encode({'messages': _wrap_positive(text)})
	feat['labels'] = [0]
	features.append(feat)
	if role == 'anchor':
	feat = template.encode({'messages': _wrap_anchor(text)})
	if feat is not None:
	feat['labels'] = [1]
	else:
	feat = template.encode({'messages': _wrap_positive(text)})
	if feat is not None:
	feat['labels'] = [0]
	if feat is not None:
	features.append(feat)

[gemini-code-assist]

Since template.encode can now return None (as updated in src/twinkle/template/base.py), executing feat['labels'] = [1] without a None check will raise a TypeError. We must guard against None before mutating feat and only append valid features.

    for t in padded:
        feat = template.encode({'messages': _wrap_anchor(t or ' ')})
        if feat is not None:
            feat['labels'] = [1]
            features.append(feat)

[gemini-code-assist]

Since template.encode can now return None (as updated in src/twinkle/template/base.py), executing feat['labels'] = [1] without a None check will raise a TypeError. We must guard against None before mutating feat and only append valid features.

Suggested change

	for t in padded:
	feat = template.encode({'messages': _wrap_anchor(t or ' ')})
	feat['labels'] = [1]
	features.append(feat)
	for t in padded:
	feat = template.encode({'messages': _wrap_anchor(t or ' ')})
	if feat is not None:
	feat['labels'] = [1]
	features.append(feat)

[gemini-code-assist]

Since template.encode can now return None (as updated in src/twinkle/template/base.py), executing feat['labels'] = [1] without a None check will raise a TypeError. We must guard against None before mutating feat and only append valid features.

Suggested change

	features.append(feat)
	out = model.forward_only(inputs=features, task='embedding', return_logits=True)
	emb = out['embeddings']
	if isinstance(emb, torch.Tensor):
	for t in padded:
	feat = template.encode({'messages': _wrap_anchor(t or ' ')})
	if feat is not None:
	feat['labels'] = [1]
	features.append(feat)

[gemini-code-assist]

Since template.encode can now return None (as updated in src/twinkle/template/base.py), executing feat['labels'] = [1] without a None check will raise a TypeError. We must guard against None before mutating feat and only append valid features.

Suggested change

	# Pad to EMB_GPUS to avoid dispatch starvation.
	pad_n = EMB_GPUS - 1
	feat = template.encode({'messages': _wrap_anchor(text)})
	if feat is not None:
	feat['labels'] = [1]

[gemini-code-assist]

If a section in Cosmopedia consists of only a single long paragraph (or does not contain \n\n), rest will be empty (''). This causes len(rest) < 256 to evaluate to True, and the entire high-quality textbook section is silently discarded. We should handle single-paragraph sections gracefully by using the title as the query and the entire body as the cot when rest is empty or too short.

Suggested change

	parts = body.split('\n\n', 1)
	first_para = parts[0].strip()
	rest = parts[1].strip() if len(parts) > 1 else ''
	if len(first_para) < 256 or len(rest) < 256:
	continue
	query = f'{title}\n\n{first_para}' if title else first_para
	out.append({
	'id': _hash_id('cosmopedia', f'{title}\n{first_para[:200]}'),
	'source': 'cosmopedia-v1',
	'query': query,
	'cot': rest,
	'response': '',
	parts = body.split('\n\n', 1)
	first_para = parts[0].strip()
	rest = parts[1].strip() if len(parts) > 1 else ''
	if len(first_para) < 256 or (rest and len(rest) < 256):
	continue
	if not rest:
	query = f'Explain {title}' if title else 'Explain the concept'
	cot = body
	else:
	query = f'{title}\n\n{first_para}' if title else first_para
	cot = rest
	out.append({
	'id': _hash_id('cosmopedia', f'{title}\n{first_para[:200]}'),
	'source': 'cosmopedia-v1',
	'query': query,
	'cot': cot,
	'response': '',
	})

[gemini-code-assist]

If self._ray_get_timeout is set to 0 or 0.0 (which is a valid timeout for non-blocking polling), 0 or timeout will evaluate to timeout because 0 is falsy in Python. We should explicitly check if _rgt is None to avoid overriding valid 0 or 0.0 timeouts.

                        _rgt = getattr(self, '_ray_get_timeout', None)
                        _rgt = _rgt if _rgt is not None else timeout

[gemini-code-assist]

There is no try...finally block around the prefetch loop. If any exception occurs during generation or evaluation (e.g., CUDA out of memory, keyboard interrupt, or API error), the prefetch_pool will not be shut down, leaving background threads hanging. We should wrap the loop in a try...finally block to ensure prefetch_pool is always shut down.

Example:

        try:
            # Prepare first batch synchronously
            cur_result = _prepare_rag_batch(batch_starts[0])
            for bi, batch_start in enumerate(batch_starts):
                ...
        finally:
            prefetch_pool.shutdown(wait=True)