"""
NeuroJenML — Dataset builder, merger, validator, and HF export.
Architecture:
- Neon stores manifests (metadata, IDs, usage tracking)
- HF stores the actual JSONL files (single source of truth)
- Training pulls JSONL from HF by path, same file that's in the dataset viewer
- No JSONL stored inline in Neon — keeps the DB lean and avoids drift
Flow: build → push to HF → store manifest → (merge) → validate → train → mark used
"""
import os
import io
import json
import time
import asyncio
from datetime import datetime
import store
# ─── JSONL builders ──────────────────────────────────────────────
# Maximum training examples generated from a single paper.
# Each map edge produces 3 examples; without a cap, papers with many edges
# (e.g. cardiovascular / MI papers with 12+ edges → 36+ examples) numerically
# dominate the training set and cause domain drift.
_MAX_EXAMPLES_PER_PAPER = 15
def _systemic_examples(extractions: list) -> list:
examples = []
for ext in extractions:
fields = ext.get("fields", {})
cls = ext.get("classification", {})
tags = cls.get("systemic_tags", [])
hypothesis = fields.get("hypothesis", "")
results = fields.get("results", "")
maps = fields.get("maps", []) or []
# Track how many examples this single paper contributes so that papers
# with many edges (e.g. cardiovascular papers with 12+ edges) cannot
# dominate the training set.
paper_example_count = 0
for m in maps:
if not (m.get("from") and m.get("to")):
continue
if paper_example_count >= _MAX_EXAMPLES_PER_PAPER:
break
relation = m.get("relation", "relates_to")
# Build a rationale that is explicitly conditioned on THIS edge
# (from → relation → to), not a raw slice of the paper body.
# A short context snippet is appended for grounding; the edge
# description is always the lead sentence so the model learns the
# correct association rather than surface co-occurrence.
context_snippet = (results or hypothesis)[:200].strip()
rationale = (
f"{m['from']} drives '{relation}', leading to {m['to']}."
+ (f" Context: {context_snippet}" if context_snippet else "")
)
# Task 1: Mechanism inference — given correlation, predict mechanism
examples.append({
"task_type": "mechanism_inference",
"instruction": (
"Given an observed association between a peripheral/systemic factor and a "
"central (brain) outcome in Alzheimer's disease, state the most likely "
"bridging mechanism. Respond as JSON with keys 'mechanism', 'rationale'."
),
"input": f"Association: '{m['from']}' is linked to '{m['to']}'. What mechanism bridges them?",
"output": json.dumps({
"mechanism": relation,
"rationale": rationale,
}),
"systemic_tags": tags,
"paper_id": ext.get("paper_id"),
"provenance": ext.get("provenance", {}),
})
# Task 2: Forward prediction — given peripheral factor, predict central outcome
examples.append({
"task_type": "forward_prediction",
"instruction": (
"You are a systemic Alzheimer's reasoning model. Given a peripheral "
"(body-level) finding, predict the downstream brain pathology and the "
"biological mechanism that links them. Respond as JSON with keys "
"'mechanism', 'central', 'rationale'."
),
"input": f"Peripheral finding: {m['from']}.",
"output": json.dumps({
"mechanism": relation,
"central": m["to"],
"rationale": rationale,
}),
"systemic_tags": tags,
"paper_id": ext.get("paper_id"),
"provenance": ext.get("provenance", {}),
})
# Task 3: Reverse prediction — given central outcome, predict peripheral cause
examples.append({
"task_type": "reverse_prediction",
"instruction": (
"Given a central brain pathology in Alzheimer's disease, identify "
"the peripheral (body-level) factors that may contribute to it. "
"Respond as JSON with keys 'peripheral', 'mechanism', 'rationale'."
),
"input": f"Central pathology: {m['to']}.",
"output": json.dumps({
"peripheral": m["from"],
"mechanism": relation,
"rationale": rationale,
}),
"systemic_tags": tags,
"paper_id": ext.get("paper_id"),
"provenance": ext.get("provenance", {}),
})
paper_example_count += 3 # 3 tasks per edge
return examples
def _reasoning_examples(extractions: list) -> list:
"""Generate chain-of-thought (CoT) reasoning training examples.
Three reasoning task types that force the model to think rather than retrieve:
reasoning_chain
A 5-step structured CoT trace wrapped in ... before the
final answer. Each step is labelled so the model learns to scaffold
its own reasoning process.
counterfactual_reasoning
"If X were absent or successfully treated, what would happen to Y?"
Forces causal-direction reasoning rather than rote association recall.
multi_hop_chain
When two edges from the same paper share a middle node (A->B, B->C),
generate a chain question: "given only A, trace through to C." This
teaches the model to compose two-step inference rather than looking up
a direct link.
"""
examples = []
for ext in extractions:
fields = ext.get("fields", {})
cls = ext.get("classification", {})
tags = cls.get("systemic_tags", [])
hypothesis = fields.get("hypothesis", "")
results = fields.get("results", "")
maps = fields.get("maps", []) or []
paper_id = ext.get("paper_id")
provenance = ext.get("provenance", {})
paper_reasoning_count = 0
# Index edges by their 'from' node so we can find chains A->B, B->C.
edge_index: dict = {}
for m in maps:
if m.get("from") and m.get("to"):
edge_index.setdefault(m["from"], []).append(
(m.get("relation", "relates_to"), m["to"])
)
for m in maps:
if not (m.get("from") and m.get("to")):
continue
if paper_reasoning_count >= _MAX_EXAMPLES_PER_PAPER:
break
src = m["from"]
tgt = m["to"]
relation = m.get("relation", "relates_to")
context_snippet = (results or hypothesis)[:200].strip()
# ── Task A: reasoning_chain ──────────────────────────────────────
# A labelled 5-step CoT inside blocks, followed by the final
# JSON answer. The model learns to narrate each inference step.
think_block = (
"\n"
f"Step 1 - Identify domain: '{src}' is a peripheral/systemic factor; "
f"'{tgt}' is a central brain outcome in Alzheimer's disease.\n"
f"Step 2 - Recall known link: prior evidence associates {src} with {tgt} "
f"through the '{relation}' pathway.\n"
f"Step 3 - Biological mechanism: {src} drives '{relation}', which in the "
f"CNS context results in {tgt}. "
+ (f"Supporting context: {context_snippet}." if context_snippet else "") + "\n"
f"Step 4 - Confidence assessment: the relationship is supported by "
f"experimental evidence from the source paper.\n"
f"Step 5 - Final answer: the mechanism is '{relation}'; the downstream "
f"central pathology is '{tgt}'.\n"
""
)
examples.append({
"task_type": "reasoning_chain",
"instruction": (
"Think step-by-step before answering. "
"Given a peripheral/systemic factor, reason through the biological "
"pathway to identify the central (brain) outcome in Alzheimer's disease. "
"First show your reasoning inside ... tags, then respond "
"with JSON: {\"mechanism\": ..., \"central\": ..., \"confidence\": ...}."
),
"input": f"Peripheral factor: {src}.",
"output": (
think_block + "\n"
+ json.dumps({
"mechanism": relation,
"central": tgt,
"confidence": "supported by experimental evidence",
})
),
"systemic_tags": tags,
"paper_id": paper_id,
"provenance": provenance,
})
paper_reasoning_count += 1
# ── Task B: counterfactual_reasoning ─────────────────────────────
# Force causal-direction reasoning: the model must reason about what
# would *not* happen if the upstream factor were removed or treated.
cf_think = (
"\n"
f"Counterfactual premise: suppose '{src}' is absent or successfully treated.\n"
f"Step 1 - Causal chain: '{src}' normally drives '{relation}', "
f"which leads to '{tgt}'.\n"
f"Step 2 - Counterfactual intervention: if '{src}' is removed, the "
f"'{relation}' signal is attenuated or eliminated.\n"
f"Step 3 - Downstream effect: without the '{relation}' driver, "
f"the manifestation of '{tgt}' would be reduced or delayed.\n"
f"Step 4 - Caveats: redundant pathways may partially compensate; "
"the effect size depends on whether this is the primary upstream driver.\n"
f"Step 5 - Conclusion: treating or removing '{src}' is predicted to "
f"attenuate '{tgt}' via the '{relation}' mechanism.\n"
""
)
examples.append({
"task_type": "counterfactual_reasoning",
"instruction": (
"Think step-by-step. "
"Given a counterfactual scenario: if a peripheral factor were absent "
"or successfully treated, reason through what would happen to the "
"downstream brain outcome. "
"Show your reasoning inside ... tags, then respond "
"with JSON: {\"effect\": ..., \"mechanism\": ..., \"caveats\": ...}."
),
"input": (
f"Counterfactual: if '{src}' were absent or successfully treated, "
f"what would happen to '{tgt}'?"
),
"output": (
cf_think + "\n"
+ json.dumps({
"effect": f"'{tgt}' would be reduced or delayed",
"mechanism": f"removal of '{src}' attenuates the '{relation}' pathway",
"caveats": "redundant upstream pathways may partially compensate",
})
),
"systemic_tags": tags,
"paper_id": paper_id,
"provenance": provenance,
})
paper_reasoning_count += 1
# ── Task C: multi_hop_chain ──────────────────────────────────────
# Find a second edge B->C where B == tgt (current edge A->B).
# Generate a two-hop chain question: "given only A, trace to C."
for next_relation, next_tgt in edge_index.get(tgt, []):
if paper_reasoning_count >= _MAX_EXAMPLES_PER_PAPER:
break
hop_think = (
"\n"
f"Multi-hop chain: '{src}' -> '{tgt}' -> '{next_tgt}'.\n"
f"Step 1 - First hop: '{src}' drives '{relation}', leading to '{tgt}'.\n"
f"Step 2 - Second hop: '{tgt}', now acting as a secondary driver, "
f"drives '{next_relation}', leading to '{next_tgt}'.\n"
f"Step 3 - Chain composition: the net effect of '{src}' on '{next_tgt}' "
f"is mediated sequentially through '{tgt}'.\n"
f"Step 4 - Direct vs. indirect: this is an indirect effect; "
f"there may be no direct single-step edge from '{src}' to '{next_tgt}'.\n"
f"Step 5 - Conclusion: '{src}' reaches '{next_tgt}' indirectly via "
f"'{tgt}', through '{relation}' then '{next_relation}'.\n"
""
)
examples.append({
"task_type": "multi_hop_chain",
"instruction": (
"Think step-by-step. "
"Given a starting peripheral factor, trace a two-step biological "
"chain to identify the final downstream brain outcome in "
"Alzheimer's disease. "
"Show your reasoning inside ... tags, then respond "
"with JSON: {\"intermediate\": ..., \"final_outcome\": ..., "
"\"step1_mechanism\": ..., \"step2_mechanism\": ...}."
),
"input": (
f"Starting factor: '{src}'. "
f"What is the downstream brain outcome after two biological steps?"
),
"output": (
hop_think + "\n"
+ json.dumps({
"intermediate": tgt,
"final_outcome": next_tgt,
"step1_mechanism": relation,
"step2_mechanism": next_relation,
})
),
"systemic_tags": tags,
"paper_id": paper_id,
"provenance": provenance,
})
paper_reasoning_count += 1
break # one multi-hop per originating edge is sufficient
return examples
def _therapeutic_reasoning_examples(extractions: list) -> list:
"""Generate therapeutic reasoning training examples.
These force the model to think like a drug designer, not a data retriever.
The model must reason about:
- What drug targets could address a mechanism
- What barriers exist to targeting a pathway
- What side effects might arise from intervention
- How to design experiments to test hypotheses
"""
import random as _rand
examples = []
paper_count = {}
for ext in extractions:
fields = ext.get("fields", {})
maps = fields.get("maps", []) or []
hypothesis = fields.get("hypothesis", "")
results = fields.get("results", "")
paper_id = ext.get("paper_id", "unknown")
if paper_count.get(paper_id, 0) >= _MAX_EXAMPLES_PER_PAPER:
continue
if not maps:
continue
for m in maps:
src = m.get("from", "")
tgt = m.get("to", "")
relation = m.get("relation", "relates_to")
if not src or not tgt:
continue
if paper_count.get(paper_id, 0) >= _MAX_EXAMPLES_PER_PAPER:
break
# Task: Therapeutic target identification
think = (
"\n"
f"The mechanism linking {src} to {tgt} in Alzheimer's disease "
f"involves '{relation}'. To intervene therapeutically, I need to consider "
f"where in this pathway a drug could act.\n\n"
f"First, what is the biological nature of this pathway? The '{relation}' "
f"mechanism connects a peripheral factor to a central brain outcome. "
f"This means there are multiple potential intervention points: "
f"reducing the peripheral driver, blocking the transport/signaling mechanism, "
f"or protecting the brain from downstream effects.\n\n"
f"What are the barriers? The blood-brain barrier limits direct CNS access. "
f"Systemic interventions may have off-target effects. The timing of intervention "
f"matters — early intervention before neurodegeneration may be more effective.\n\n"
f"What drug modalities could work? Antibodies targeting the peripheral factor, "
f"small molecules blocking the signaling pathway, or gene therapies modifying "
f"the expression of key mediators.\n\n"
f"The most promising approach depends on the specificity of the '{relation}' "
f"mechanism and whether it is a primary driver or a compensatory pathway.\n"
""
)
examples.append({
"task_type": "therapeutic_reasoning",
"instruction": (
"You are a neuropharmacologist designing a therapeutic strategy for "
"Alzheimer's disease. Analyze the given mechanism and reason through "
"potential drug targets, delivery challenges, and experimental approaches. "
"Show your reasoning in ..., then respond with analysis."
),
"input": (
f"Mechanism: {src} drives '{relation}', leading to {tgt}. "
f"What are the therapeutic implications and drug design considerations?"
),
"output": (
think + "\n"
+ f"The therapeutic strategy for targeting the '{relation}' pathway "
f"between {src} and {tgt} should focus on: (1) identifying the most "
f"druggable node in the pathway, (2) developing delivery mechanisms "
f"that cross the blood-brain barrier where needed, and "
f"(3) designing biomarkers to monitor treatment response. "
f"The key challenge is that systemic interventions targeting {src} "
f"may have unintended consequences on other physiological processes."
),
"paper_id": paper_id,
})
paper_count[paper_id] = paper_count.get(paper_id, 0) + 1
# Task: Experimental design reasoning
think2 = (
"\n"
f"To test whether {src} causally contributes to {tgt} via '{relation}', "
f"I need to design an experiment that isolates this specific pathway.\n\n"
f"A correlation study alone is insufficient — I need interventional evidence. "
f"Options include: (1) a knockout model where the {src} pathway is disabled, "
f"(2) an antibody blockade study reducing circulating {src}, or "
f"(3) a longitudinal cohort tracking {src} levels and {tgt} progression.\n\n"
f"Each approach has tradeoffs. Knockout models prove causality but may not "
f"translate to human therapy. Antibody blockade is more translatable but "
f"requires knowing the right timing window. Longitudinal studies establish "
f"temporal relationships but cannot prove mechanism.\n\n"
f"The strongest evidence would combine approaches: use the longitudinal data "
f"to identify the critical time window, then use the interventional studies "
f"to test causality within that window.\n"
""
)
examples.append({
"task_type": "experimental_design",
"instruction": (
"Design an experimental approach to test whether a peripheral factor "
"causally contributes to brain pathology in Alzheimer's disease. "
"Reason through the experimental design in ..., "
"then describe your recommended approach."
),
"input": (
f"Hypothesis: {src} contributes to {tgt} via '{relation}'. "
f"Design an experiment to test this."
),
"output": (
think2 + "\n"
+ f"To test the causal link between {src} and {tgt}, I recommend "
f"a multi-pronged approach: first establish the temporal relationship "
f"through longitudinal biomarker tracking, then use targeted intervention "
f"(antibody blockade or pathway inhibition) to test whether reducing "
f"the {src} signal attenuates {tgt} progression. Controls should include "
f"pathway-specific inhibitors and sham-treated animals."
),
"paper_id": paper_id,
})
paper_count[paper_id] = paper_count.get(paper_id, 0) + 1
return examples
def _general_examples(records: list) -> list:
examples = []
for rec in records:
params = rec.get("parameters", {})
if not params:
continue
examples.append({
"task_type": "general_context",
"instruction": (
"Given standard demographic/clinical parameters, assess Alzheimer's disease "
"risk context. Respond as JSON with key 'assessment'."
),
"input": json.dumps(params),
"output": json.dumps({"assessment": "", "species_class": rec.get("species_class")}),
"species_class": rec.get("species_class"),
"dataset_name": rec.get("dataset_name"),
})
return examples
def _maps_examples(maps: list) -> list:
examples = []
for m in maps:
if not (m.get("from") and m.get("to")):
continue
examples.append({
"task_type": "map_forward",
"instruction": "Predict the downstream node. Respond as JSON with key 'to'.",
"input": json.dumps({"from": m["from"], "relation": m.get("relation")}),
"output": json.dumps({"to": m["to"]}),
"domain": m.get("domain", ""),
})
examples.append({
"task_type": "map_reverse",
"instruction": "Reverse-engineer the upstream node. Respond as JSON with key 'from'.",
"input": json.dumps({"to": m["to"], "relation": m.get("relation")}),
"output": json.dumps({"from": m["from"]}),
"domain": m.get("domain", ""),
})
return examples
async def _gather_examples(kind: str) -> list:
if kind == "systemic":
extractions = await store.list_all("extractions", newest_first=True)
return _systemic_examples(extractions)
if kind == "general":
records = await store.list_all("general_records", newest_first=True)
return _general_examples(records)
if kind == "maps":
maps = await store.list_all("maps", newest_first=True)
return _maps_examples(maps)
if kind == "reasoning":
# Reasoning examples are derived from the same extractions as systemic
# examples. They deliberately overlap in source data but teach a
# completely different cognitive skill: structured inference rather than
# pattern-matched retrieval.
extractions = await store.list_all("extractions", newest_first=True)
return _reasoning_examples(extractions)
if kind == "therapeutic":
# Therapeutic reasoning teaches the model to think like a drug designer:
# analyzing mechanisms, identifying targets, designing experiments.
extractions = await store.list_all("extractions", newest_first=True)
return _therapeutic_reasoning_examples(extractions)
if kind == "all_reasoning":
# Combined reasoning + therapeutic examples for comprehensive training.
extractions = await store.list_all("extractions", newest_first=True)
return _reasoning_examples(extractions) + _therapeutic_reasoning_examples(extractions)
return []
# ─── Normalization ───────────────────────────────────────────────
_CANONICAL_KEYS = ("task_type", "instruction", "input", "output",
"systemic_tags", "paper_id", "provenance",
"species_class", "dataset_name", "domain",
# reasoning examples carry their CoT inside output; no extra keys needed
)
def _normalize_example(ex: dict) -> dict:
out = {}
for k in _CANONICAL_KEYS:
out[k] = ex.get(k, "" if k != "systemic_tags" else [])
return out
def _to_jsonl(examples: list) -> str:
return "\n".join(json.dumps(_normalize_example(ex), ensure_ascii=False) for ex in examples)
def parse_jsonl(jsonl: str) -> list:
examples = []
for line in jsonl.strip().split("\n"):
line = line.strip()
if not line:
continue
try:
obj = json.loads(line)
if isinstance(obj, dict):
examples.append(obj)
except (json.JSONDecodeError, ValueError):
pass
return examples
def merge_jsonl_parts(parts: list) -> str:
"""Merge multiple JSONL strings, deduplicating by (task_type, input) pair."""
all_examples = []
for part in parts:
all_examples.extend(parse_jsonl(part))
seen = set()
unique = []
for ex in all_examples:
key = (ex.get("task_type", ""), ex.get("input", ""))
if key not in seen:
seen.add(key)
unique.append(ex)
return _to_jsonl(unique)
# ─── HF Hub operations ───────────────────────────────────────────
def _hf_token():
return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
def _hf_repo():
return os.getenv("HF_DATASET_REPO")
def _push_to_hf(kind: str, version: str, jsonl: str) -> dict:
token = _hf_token()
repo = _hf_repo()
if not token or not repo:
return {"pushed": False, "reason": "HF_TOKEN or HF_DATASET_REPO not set"}
try:
from huggingface_hub import HfApi
api = HfApi(token=token)
api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
path_in_repo = f"{kind}/{version}/training_data.jsonl"
api.upload_file(
path_or_fileobj=io.BytesIO(jsonl.encode("utf-8")),
path_in_repo=path_in_repo,
repo_id=repo,
repo_type="dataset",
commit_message=f"NeuroJenML {kind} dataset {version}",
)
return {"pushed": True, "repo": repo, "path": path_in_repo}
except Exception as e:
return {"pushed": False, "reason": str(e)[:300]}
def pull_from_hf(path_in_repo: str, max_retries: int = 3) -> str:
"""Pull a JSONL file from HF dataset repo by path. Retries on transient failures."""
token = _hf_token()
repo = _hf_repo()
if not token or not repo:
raise RuntimeError("HF_TOKEN or HF_DATASET_REPO not set")
from huggingface_hub import hf_hub_download
last_error = None
for attempt in range(max_retries):
try:
local_path = hf_hub_download(
repo_id=repo,
filename=path_in_repo,
repo_type="dataset",
token=token,
)
with open(local_path, "r", encoding="utf-8") as f:
return f.read()
except Exception as e:
last_error = e
if attempt < max_retries - 1:
import time as _time
delay = 2 ** attempt
print(f"[retry] HF pull attempt {attempt + 1}/{max_retries} failed: {str(e)[:100]}. Retrying in {delay}s...")
_time.sleep(delay)
raise last_error
def push_extraction_to_hf(paper_id: str, extraction: dict, repo: str, token: str) -> str:
try:
from huggingface_hub import HfApi
except ImportError:
raise RuntimeError("huggingface_hub not installed")
import json as _json
record = {
"paper_id": paper_id,
"model": extraction.get("model"),
"classification": extraction.get("classification"),
"fields": extraction.get("fields"),
"provenance": extraction.get("provenance"),
"extracted_at": datetime.now().isoformat(),
}
content = (_json.dumps(record, ensure_ascii=False) + "\n").encode("utf-8")
path_in_repo = f"extractions/{paper_id}.jsonl"
api = HfApi(token=token)
api.upload_file(
path_or_fileobj=content,
path_in_repo=path_in_repo,
repo_id=repo,
repo_type="dataset",
commit_message=f"add extraction {paper_id}",
)
return f"https://huggingface.co/datasets/{repo}/blob/main/{path_in_repo}"
# ─── Build & Store (manifest only, JSONL goes to HF) ─────────────
async def build_and_record(kind: str, push: bool = True) -> dict:
"""Build a dataset, push JSONL to HF, store manifest in Neon."""
examples = await _gather_examples(kind)
version = "v" + datetime.now().strftime("%Y%m%d-%H%M%S")
jsonl = _to_jsonl(examples)
push_meta = {"pushed": False}
if push:
push_meta = _push_to_hf(kind, version, jsonl)
ds_id = f"ds-{kind}-{int(time.time() * 1000) % 10_000_000}"
manifest = {
"id": ds_id,
"kind": kind,
"version": version,
"example_count": len(examples),
"hf_path": push_meta.get("path", ""),
"hf_repo": push_meta.get("repo", ""),
"hf_pushed": push_meta.get("pushed", False),
# Inline JSONL stored as fallback when HF push is unavailable.
# This is what the training bridge uses when hf_path is empty.
# Capped at 8 MB to stay within Neon JSONB limits.
"jsonl_inline": jsonl if len(jsonl) < 8 * 1024 * 1024 else "",
"source_papers": list({ex.get("paper_id") for ex in examples if ex.get("paper_id")}),
"task_types": list({ex.get("task_type") for ex in examples if ex.get("task_type")}),
"used_in_training": False,
"training_job_id": None,
"created_at": datetime.now().isoformat(),
}
await store.upsert("datasets", ds_id, manifest)
return {
"status": "ok",
"dataset": manifest,
"jsonl": jsonl,
"count": len(examples),
}
# ─── Dataset Operations ──────────────────────────────────────────
async def list_datasets(kind: str = None, unused_only: bool = False) -> list:
filters = {}
if kind:
filters["kind"] = kind
datasets = await store.query("datasets", filters, newest_first=True)
if unused_only:
datasets = [d for d in datasets if not d.get("used_in_training")]
return datasets
async def get_dataset(ds_id: str) -> dict:
return await store.get("datasets", ds_id) or {}
async def get_dataset_jsonl(ds_id: str) -> str:
"""Return the JSONL for a dataset, pulling from HF or falling back to inline.
Priority:
1. HF Hub (hf_path set and HF credentials available)
2. Inline JSONL stored in the manifest at build time (jsonl_inline field)
"""
ds = await store.get("datasets", ds_id)
if not ds:
return ""
hf_path = ds.get("hf_path", "")
if hf_path:
try:
return await asyncio.to_thread(pull_from_hf, hf_path)
except Exception:
pass # fall through to inline
return ds.get("jsonl_inline", "")
async def merge_datasets(ds_ids: list, name: str = None) -> dict:
"""Merge multiple datasets by pulling each from HF and combining."""
all_examples = []
source_ids = []
source_papers = set()
for ds_id in ds_ids:
ds = await store.get("datasets", ds_id)
if not ds:
continue
hf_path = ds.get("hf_path", "")
if not hf_path:
continue # Skip datasets without HF path
try:
jsonl = await asyncio.to_thread(pull_from_hf, hf_path)
except Exception:
jsonl = ""
if not jsonl:
continue
examples = parse_jsonl(jsonl)
all_examples.extend(examples)
source_ids.append(ds_id)
source_papers.update(ds.get("source_papers", []))
# Deduplicate by (task_type, input)
seen = set()
unique = []
for ex in all_examples:
key = (ex.get("task_type", ""), ex.get("input", ""))
if key not in seen:
seen.add(key)
unique.append(ex)
version = "v" + datetime.now().strftime("%Y%m%d-%H%M%S")
jsonl = _to_jsonl(unique)
# Push merged dataset to HF
push_meta = _push_to_hf("merged", version, jsonl)
ds_id = f"ds-merged-{int(time.time() * 1000) % 10_000_000}"
manifest = {
"id": ds_id,
"kind": "merged",
"version": version,
"name": name or f"Merged from {len(source_ids)} datasets",
"example_count": len(unique),
"hf_path": push_meta.get("path", ""),
"hf_repo": push_meta.get("repo", ""),
"hf_pushed": push_meta.get("pushed", False),
"source_datasets": source_ids,
"source_papers": list(source_papers),
"task_types": list({ex.get("task_type") for ex in unique if ex.get("task_type")}),
"jsonl_inline": jsonl if len(jsonl) < 8 * 1024 * 1024 else "",
"used_in_training": False,
"training_job_id": None,
"created_at": datetime.now().isoformat(),
}
await store.upsert("datasets", ds_id, manifest)
return {
"status": "ok",
"dataset": manifest,
"count": len(unique),
"deduplicated": len(all_examples) - len(unique),
}
async def validate_dataset(ds_id: str) -> dict:
"""Validate a dataset pulled from HF. Checks structure and content quality."""
try:
ds = await store.get("datasets", ds_id)
jsonl = await asyncio.to_thread(
pull_from_hf, (ds or {}).get("hf_path", "")
)
except Exception as e:
return {"valid": False, "error": f"Could not pull from HF: {str(e)[:200]}"}
if not jsonl or not jsonl.strip():
return {"valid": False, "error": "Dataset is empty"}
examples = parse_jsonl(jsonl)
issues = []
required_keys = {"task_type", "instruction", "input", "output"}
seen_pairs = set()
leakage_count = 0
empty_fields = 0
for i, ex in enumerate(examples):
missing = required_keys - set(ex.keys())
if missing:
issues.append(f"Line {i + 1}: missing keys {missing}")
for key in ["instruction", "input", "output"]:
val = ex.get(key, "")
if not val or (isinstance(val, str) and not val.strip()):
empty_fields += 1
pair = (ex.get("task_type", ""), ex.get("input", ""))
if pair in seen_pairs:
issues.append(f"Line {i + 1}: duplicate (task_type, input) pair")
seen_pairs.add(pair)
output_val = ex.get("output", "")
input_val = ex.get("input", "")
if isinstance(output_val, str) and isinstance(input_val, str):
try:
out_obj = json.loads(output_val)
for v in out_obj.values():
if isinstance(v, str) and len(v) > 5 and v.lower() in input_val.lower():
leakage_count += 1
break
except (json.JSONDecodeError, ValueError):
pass
total = len(examples)
quality_score = 100.0
if total == 0:
quality_score = 0
else:
quality_score -= (empty_fields / max(1, total * 3)) * 30
quality_score -= (leakage_count / total) * 40
quality_score -= min(20, len(issues) * 2)
quality_score = max(0, round(quality_score, 1))
return {
"valid": len(issues) == 0 and total > 0,
"example_count": total,
"quality_score": quality_score,
"issues": issues[:20],
"issue_count": len(issues),
"empty_fields": empty_fields,
"leakage_examples": leakage_count,
"duplicate_pairs": len(examples) - len(seen_pairs),
}
async def mark_used(ds_id: str, job_id: str) -> bool:
ds = await store.get("datasets", ds_id)
if not ds:
return False
ds["used_in_training"] = True
ds["training_job_id"] = job_id
ds["used_at"] = datetime.now().isoformat()
await store.upsert("datasets", ds_id, ds)
return True
async def mark_unused(ds_id: str) -> bool:
"""Allow a used dataset to be reused for training."""
ds = await store.get("datasets", ds_id)
if not ds:
return False
ds["used_in_training"] = False
ds["reused"] = True
ds["reused_at"] = datetime.now().isoformat()
await store.upsert("datasets", ds_id, ds)
return True