| """Evidence-based agentic dataset discovery.""" |
| from __future__ import annotations |
|
|
| import json |
| import math |
| import os |
| import re |
| import threading |
| import time |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from typing import Any, Iterator |
|
|
| from backend.search import inspect_dataset, search_datasets |
|
|
| MODEL = os.getenv("WEAVER_MODEL", "HuggingFaceTB/SmolLM2-360M-Instruct") |
| MAX_TASK_LENGTH = 2000 |
|
|
| _local_model = None |
| _local_tokenizer = None |
| _model_load_failed = False |
| _model_lock = threading.Lock() |
| _WORD_RE = re.compile(r"[a-zA-Z][a-zA-Z0-9_-]{2,}") |
| _STOPWORDS = { |
| "the", "and", "for", "with", "that", "this", "from", "your", "have", |
| "need", "want", "like", "build", "using", "data", "dataset", "model", |
| "small", "find", "looking", "project", "create", "make", "into", |
| "about", "pairs", "examples", "records", "documents", "corpus", |
| "evaluation", "evaluate", "permissive", "license", "compact", |
| "abstractive", "extractive", "recordings", "transcripts", |
| } |
| _LANGUAGES = { |
| "english": "en", "arabic": "ar", "french": "fr", "german": "de", |
| "spanish": "es", "italian": "it", "portuguese": "pt", "chinese": "zh", |
| "japanese": "ja", "korean": "ko", "hindi": "hi", "multilingual": "multilingual", |
| } |
| _MODALITIES = { |
| "audio": ("audio", "speech", "voice", "asr"), |
| "image": ("image", "vision", "photo", "ocr"), |
| "video": ("video",), |
| "tabular": ("tabular", "table", "csv", "structured"), |
| "text": ("text", "document", "summarization", "translation", "intent", "chat"), |
| } |
| _LABEL_TERMS = {"label", "labels", "class", "classes", "intent", "category"} |
| _DIRECT_LABEL_FIELDS = { |
| "label", "labels", "intent", "category", "class", "target", |
| "detected_intent", "intent_label", "class_label", |
| } |
| _PROXY_LABEL_FIELDS = {"type", "queue", "topic", "department", "route", "routing"} |
| _TASK_TYPES = { |
| "intent classification", "classification", "summarization", "translation", |
| "question answering", "retrieval", "automatic speech recognition", |
| "fine-tuning", "pretraining", "dataset discovery", |
| } |
| _TASK_ALIASES = { |
| "intent classification": ("intent classification", "intent"), |
| "classification": ("classification",), |
| "summarization": ("summarization", "summary"), |
| "translation": ("translation",), |
| "question answering": ("question answering", "qa"), |
| "retrieval": ("retrieval", "search"), |
| "automatic speech recognition": ("speech recognition", "asr", "transcription"), |
| "fine-tuning": ("instruction", "fine tuning"), |
| "pretraining": ("pretraining",), |
| } |
| _FIELD_ALIASES = { |
| "text": { |
| "text", "sentence", "content", "document", "article", "body", "query", |
| "utterance", "input_text", "text_input", |
| }, |
| "label": _DIRECT_LABEL_FIELDS, |
| "document": { |
| "document", "article", "text", "content", "body", "source", |
| "judgement", "judgment", "case_text", "legal_text", |
| }, |
| "summary": { |
| "summary", "highlights", "abstract", "target", "headline", |
| "summarizer", "processed_text", |
| }, |
| "question": {"question", "query", "prompt", "question_text", "instruction"}, |
| "answer": {"answer", "answers", "response", "context", "answer_text"}, |
| "source": {"source", "src", "text", "sentence", "input", "source_text", "input_text"}, |
| "target": { |
| "target", "tgt", "translation", "translated_text", "output", |
| "target_text", "output_text", |
| }, |
| "audio": {"audio", "speech", "file", "path", "audio_path", "audio_file"}, |
| "transcript": { |
| "transcript", "transcription", "sentence", "text", |
| "transcript_text", "transcription_text", "label", |
| }, |
| "instruction": {"instruction", "prompt", "input"}, |
| "response": {"response", "output", "completion", "answer"}, |
| } |
|
|
|
|
| def _domain_terms(profile: dict[str, Any]) -> list[str]: |
| task_words = { |
| word |
| for alias in _TASK_ALIASES.get(profile["task_type"], (profile["task_type"],)) |
| for word in alias.split() |
| } |
| ignored = { |
| *profile["languages"], |
| *_LANGUAGES.keys(), |
| *task_words, |
| *profile["required_fields"], |
| "labels", "label", "classifier", "classification", |
| "summarization", "summary", "translation", "retrieval", "search", |
| "question", "answer", "speech", "recognition", "asr", "transcript", |
| "text", "audio", "image", "video", |
| } |
| return [term for term in profile["domain_keywords"] if term not in ignored] |
|
|
|
|
| def _llm(system: str, user: str, max_tokens: int = 350, temperature: float = 0.2) -> str | None: |
| """Run the planning step locally with a genuinely small model.""" |
| global _local_model, _local_tokenizer, _model_load_failed |
| if _model_load_failed: |
| return None |
| if _local_model is None or _local_tokenizer is None: |
| try: |
| with _model_lock: |
| if _local_model is None or _local_tokenizer is None: |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| _local_tokenizer = AutoTokenizer.from_pretrained(MODEL) |
| _local_model = AutoModelForCausalLM.from_pretrained( |
| MODEL, |
| low_cpu_mem_usage=True, |
| ) |
| _local_model.eval() |
| except Exception: |
| _model_load_failed = True |
| return None |
| try: |
| messages = [ |
| {"role": "system", "content": system}, |
| {"role": "user", "content": user}, |
| ] |
| prompt = _local_tokenizer.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
| inputs = _local_tokenizer(prompt, return_tensors="pt") |
| generation_args = { |
| "max_new_tokens": min(max_tokens, 220), |
| "do_sample": temperature > 0, |
| "pad_token_id": _local_tokenizer.eos_token_id, |
| } |
| if temperature > 0: |
| generation_args["temperature"] = temperature |
| with _model_lock: |
| output = _local_model.generate( |
| **inputs, |
| **generation_args, |
| ) |
| generated = output[0, inputs["input_ids"].shape[1]:] |
| return _local_tokenizer.decode( |
| generated, |
| skip_special_tokens=True, |
| ) |
| except Exception: |
| return None |
|
|
|
|
| def _extract_json(text: str | None) -> dict[str, Any]: |
| if not text: |
| return {} |
| cleaned = re.sub(r"^```(?:json)?\s*|\s*```$", "", text.strip(), flags=re.I) |
| try: |
| parsed = json.loads(cleaned) |
| return parsed if isinstance(parsed, dict) else {} |
| except Exception: |
| pass |
| start = cleaned.find("{") |
| end = cleaned.rfind("}") |
| if start >= 0 and end > start: |
| try: |
| parsed = json.loads(cleaned[start:end + 1]) |
| return parsed if isinstance(parsed, dict) else {} |
| except Exception: |
| pass |
| return {} |
|
|
|
|
| def _keywords(text: str) -> list[str]: |
| normalized = text.replace("-", " ") |
| return [word.lower() for word in _WORD_RE.findall(normalized) if word.lower() not in _STOPWORDS] |
|
|
|
|
| def _task_type(lower: str) -> str: |
| if "intent" in lower and any(term in lower for term in ("classif", "label", "dataset", "data")): |
| return "intent classification" |
| if "summar" in lower: |
| return "summarization" |
| if "translat" in lower: |
| return "translation" |
| if "retrieval" in lower or "search evaluation" in lower: |
| return "retrieval" |
| if "question answer" in lower or re.search(r"\bqa\b", lower): |
| return "question answering" |
| if re.search(r"\basr\b", lower) or "speech recognition" in lower or "speech to text" in lower: |
| return "automatic speech recognition" |
| if "classif" in lower or "classifier" in lower: |
| return "classification" |
| if "fine-tun" in lower or "finetun" in lower: |
| return "fine-tuning" |
| if "pretrain" in lower: |
| return "pretraining" |
| return "dataset discovery" |
|
|
|
|
| def _default_required_fields(task_type: str) -> list[str]: |
| return { |
| "intent classification": ["text", "label"], |
| "classification": ["text", "label"], |
| "summarization": ["document", "summary"], |
| "translation": ["source", "target"], |
| "question answering": ["question", "answer"], |
| "automatic speech recognition": ["audio", "transcript"], |
| }.get(task_type, []) |
|
|
|
|
| def parse_task(task: str, use_llm: bool = True) -> tuple[dict[str, Any], bool]: |
| """Turn a free-form project description into explicit search requirements.""" |
| lower = task.lower() |
| languages = [code for name, code in _LANGUAGES.items() if name in lower] |
| modalities = [ |
| modality |
| for modality, terms in _MODALITIES.items() |
| if any(term in lower for term in terms) |
| ] |
| if not modalities: |
| modalities = ["text"] |
| task_type = _task_type(lower) |
| required_fields = _default_required_fields(task_type) |
| word_set = set(_keywords(task)) |
| if task_type not in {"translation", "summarization", "question answering"}: |
| if word_set & _LABEL_TERMS: |
| required_fields.append("label") |
| for field in ("question", "answer", "instruction", "response", "summary", "transcript"): |
| if field in word_set: |
| required_fields.append(field) |
| profile: dict[str, Any] = { |
| "languages": languages, |
| "modalities": modalities, |
| "task_type": task_type, |
| "required_fields": list(dict.fromkeys(required_fields)), |
| "license": "permissive" if any( |
| term in lower for term in ("commercial", "production", "permissive", "apache", "mit") |
| ) else "", |
| "size_preference": "small" if any( |
| term in lower for term in ("small", "tiny", "prototype", "quick") |
| ) else "", |
| "domain_keywords": _keywords(task)[:12], |
| "intended_use": task[:280], |
| } |
|
|
| llm_used = False |
| if use_llm: |
| prompt = ( |
| f"Request: {task}\n" |
| "Return exactly one compact JSON object. Use at most 5 domain_keywords. " |
| "Use an empty string or empty list when a requirement is not explicit. " |
| 'Schema: {"languages":[],"modalities":[],"task_type":"","required_fields":[],' |
| '"license":"","size_preference":"","domain_keywords":[]}' |
| ) |
| parsed = _extract_json(_llm( |
| "Extract only requirements explicitly stated. No prose. No repetition.", |
| prompt, |
| max_tokens=120, |
| temperature=0, |
| )) |
| if parsed: |
| llm_used = True |
| proposed_languages = [ |
| str(item).lower() for item in parsed.get("languages", []) |
| if isinstance(item, str) |
| and (len(item.strip()) in {2, 3} or item.lower() == "multilingual") |
| ] |
| proposed_modalities = [ |
| str(item).lower() for item in parsed.get("modalities", []) |
| if str(item).lower() in _MODALITIES |
| ] |
| proposed_fields = [ |
| str(item).lower() for item in parsed.get("required_fields", []) |
| if isinstance(item, str) and len(item) <= 30 |
| ] |
| proposed_keywords = [ |
| str(item).lower() for item in parsed.get("domain_keywords", []) |
| if isinstance(item, str) and 2 < len(item) <= 30 |
| ] |
| profile["languages"] = list(dict.fromkeys(profile["languages"] + proposed_languages)) |
| profile["modalities"] = list(dict.fromkeys(profile["modalities"] + proposed_modalities)) |
| profile["required_fields"] = list( |
| dict.fromkeys(profile["required_fields"] + proposed_fields) |
| ) |
| profile["domain_keywords"] = list( |
| dict.fromkeys(profile["domain_keywords"] + proposed_keywords) |
| )[:12] |
| proposed_task = str(parsed.get("task_type") or "").lower() |
| if profile["task_type"] == "dataset discovery" and proposed_task in _TASK_TYPES: |
| profile["task_type"] = proposed_task |
| proposed_license = str(parsed.get("license") or "").lower() |
| if profile["license"] == "" and proposed_license and proposed_license in lower: |
| profile["license"] = proposed_license |
| proposed_size = str(parsed.get("size_preference") or "").lower() |
| if profile["size_preference"] == "" and proposed_size in {"small", "medium", "large"}: |
| if proposed_size in lower: |
| profile["size_preference"] = proposed_size |
| return profile, llm_used |
|
|
|
|
| def generate_queries(task: str, profile: dict[str, Any]) -> list[str]: |
| task_type = profile["task_type"] |
| task_aliases = _TASK_ALIASES.get(task_type, (task_type,)) |
| terms = _domain_terms(profile) |
| primary_task = task_aliases[0] |
| compact_task = task_aliases[-1] |
| language_names = [ |
| name for name, code in _LANGUAGES.items() |
| if code in profile["languages"] and name != "multilingual" |
| ] |
| domain = terms[:2] |
| field_terms = [field for field in profile["required_fields"] if field not in {"text", "document", "source", "target"}] |
| queries = [] |
| if domain: |
| task_for_domain = primary_task if len(domain) + len(primary_task.split()) <= 3 else compact_task |
| queries.append(" ".join(domain + [task_for_domain])) |
| queries.append(" ".join([domain[0], compact_task])) |
| queries.append(" ".join(domain + ["dataset"])) |
| if language_names: |
| queries.append(" ".join([language_names[0], compact_task])) |
| if language_names and domain: |
| queries.append(" ".join([language_names[0], domain[0]])) |
| queries.append(" ".join([language_names[0], domain[0], "dataset"])) |
| elif len(language_names) >= 2: |
| queries.append(" ".join(language_names[:2] + [compact_task])) |
| if domain and {"question", "answer"}.issubset(profile["required_fields"]): |
| queries.append(f"{domain[0]} question") |
| queries.append(f"{domain[0]} qa") |
| queries.append(" ".join(domain + ["qa"])) |
| if domain and field_terms: |
| queries.append(" ".join(domain[:1] + field_terms[:2])) |
| if task_type == "automatic speech recognition": |
| queries.append("speech transcription") |
| queries.append("librispeech") |
| if task_type == "intent classification": |
| queries.append("intent dataset") |
| if domain: |
| queries.append(f"{domain[0]} intent") |
| if domain: |
| queries.append(" ".join(domain)) |
| queries.append(primary_task) |
| if len(task_aliases) > 1: |
| queries.append(task_aliases[-1]) |
| cleaned = [] |
| for query in queries: |
| normalized = re.sub(r"\s+", " ", query).strip() |
| if normalized and normalized not in cleaned: |
| cleaned.append(normalized) |
| return cleaned[:9] or [" ".join(_keywords(task)[:4])] |
|
|
|
|
| def _text_blob(dataset: dict[str, Any]) -> str: |
| values = [ |
| dataset.get("id", ""), |
| dataset.get("description", ""), |
| " ".join(dataset.get("tags", [])), |
| " ".join(dataset.get("features", [])), |
| " ".join(dataset.get("task_categories", [])), |
| ] |
| return " ".join(values).lower() |
|
|
|
|
| def _pre_score(profile: dict[str, Any], dataset: dict[str, Any]) -> float: |
| blob = _text_blob(dataset) |
| keywords = set(_domain_terms(profile) or profile["domain_keywords"]) |
| overlap = sum(1 for word in keywords if word in blob) |
| modality_matches = sum(1 for value in profile["modalities"] if value in dataset.get("modalities", [])) |
| language_matches = sum(1 for value in profile["languages"] if value in dataset.get("languages", [])) |
| task_terms = _TASK_ALIASES.get(profile["task_type"], (profile["task_type"],)) |
| task_match = sum(1 for term in task_terms if term in blob) |
| schema_hint = sum(1 for field in profile["required_fields"] if field in blob) |
| popularity = min(1.5, math.log10(1 + dataset.get("downloads", 0) + dataset.get("likes", 0) * 10) / 2) |
| return overlap * 10 + task_match * 10 + schema_hint * 4 + modality_matches * 8 + language_matches * 8 + popularity |
|
|
|
|
| def _contains_any(values: list[str], expected: list[str]) -> bool: |
| lowered = {str(value).lower() for value in values} |
| return any(item.lower() in lowered for item in expected) |
|
|
|
|
| def _field_names(dataset: dict[str, Any]) -> set[str]: |
| names = {str(field).lower() for field in dataset.get("features", [])} |
|
|
| def visit(value: Any) -> None: |
| if isinstance(value, dict): |
| for key, nested in value.items(): |
| names.add(str(key).lower()) |
| visit(nested) |
| elif isinstance(value, list): |
| for nested in value[:5]: |
| visit(nested) |
|
|
| for row in dataset.get("sample_rows", []): |
| visit(row) |
| return names |
|
|
|
|
| def _sample_text(dataset: dict[str, Any]) -> str: |
| values: list[str] = [] |
|
|
| def visit(value: Any) -> None: |
| if isinstance(value, dict): |
| for nested in value.values(): |
| visit(nested) |
| elif isinstance(value, list): |
| for nested in value[:5]: |
| visit(nested) |
| elif isinstance(value, str): |
| values.append(value) |
|
|
| for row in dataset.get("sample_rows", []): |
| visit(row) |
| return " ".join(values).lower() |
|
|
|
|
| def _infer_script_languages(text: str) -> list[str]: |
| if not text: |
| return [] |
| letters = [char for char in text if char.isalpha()] |
| if not letters: |
| return [] |
| ranges = { |
| "ar": lambda char: "\u0600" <= char <= "\u06ff", |
| "zh": lambda char: "\u4e00" <= char <= "\u9fff", |
| "ja": lambda char: "\u3040" <= char <= "\u30ff", |
| "ko": lambda char: "\uac00" <= char <= "\ud7af", |
| } |
| return [ |
| language |
| for language, matcher in ranges.items() |
| if sum(1 for char in letters if matcher(char)) / len(letters) >= 0.05 |
| ] |
|
|
|
|
| def _matches_field(requirement: str, field: str) -> bool: |
| aliases = _FIELD_ALIASES.get(requirement, {requirement}) |
| normalized = field.replace("-", "_").lower() |
| return normalized in aliases |
|
|
|
|
| def _field_value(row: dict[str, Any], fields: list[str]) -> Any: |
| lowered = {str(key).lower(): value for key, value in row.items()} |
| for field in fields: |
| value = lowered.get(str(field).lower()) |
| if value not in (None, "", [], {}): |
| return value |
| return None |
|
|
|
|
| def _sample_tests( |
| profile: dict[str, Any], |
| dataset: dict[str, Any], |
| matched_requirements: dict[str, list[str]], |
| ) -> list[dict[str, str]]: |
| rows = [row for row in dataset.get("sample_rows", []) if isinstance(row, dict)] |
| required_fields = profile["required_fields"] |
| tests: list[dict[str, str]] = [] |
| if not rows: |
| return [{ |
| "name": "Sample rows available", |
| "status": "unknown", |
| "detail": "Dataset Viewer did not expose sample rows.", |
| }] |
|
|
| missing = [ |
| requirement for requirement in required_fields |
| if not matched_requirements.get(requirement) |
| ] |
| empty = [ |
| requirement for requirement in required_fields |
| if matched_requirements.get(requirement) |
| and not any(_field_value(row, matched_requirements[requirement]) is not None for row in rows) |
| ] |
| tests.append({ |
| "name": "Required fields populated", |
| "status": "pass" if not missing and not empty else "fail" if missing else "review", |
| "detail": "All requested fields have sample values." |
| if not missing and not empty |
| else f"Missing: {', '.join(missing or empty)}.", |
| }) |
|
|
| if {"question", "answer"} & set(required_fields): |
| question_fields = matched_requirements.get("question", []) |
| answer_fields = matched_requirements.get("answer", []) |
| question_values = [ |
| str(_field_value(row, question_fields) or "") for row in rows |
| if _field_value(row, question_fields) is not None |
| ] |
| answer_values = [ |
| str(_field_value(row, answer_fields) or "") for row in rows |
| if _field_value(row, answer_fields) is not None |
| ] |
| plausible_questions = sum( |
| 1 for value in question_values |
| if value.endswith("?") or len(value.split()) >= 3 |
| ) |
| distinct_answers = sum( |
| 1 for question, answer in zip(question_values, answer_values) |
| if answer and answer.strip().lower() != question.strip().lower() |
| ) |
| tests.append({ |
| "name": "QA row shape", |
| "status": "pass" if plausible_questions and distinct_answers else "review", |
| "detail": "Questions and answers look usable in inspected rows." |
| if plausible_questions and distinct_answers |
| else "Question/answer shape needs manual review.", |
| }) |
|
|
| if "label" in required_fields: |
| label_fields = matched_requirements.get("label", []) |
| label_values = [ |
| _field_value(row, label_fields) for row in rows |
| if _field_value(row, label_fields) is not None |
| ] |
| unique_labels = {str(value) for value in label_values} |
| tests.append({ |
| "name": "Label signal", |
| "status": "pass" if label_values and len(unique_labels) <= max(20, len(rows)) else "review", |
| "detail": f"Found {len(unique_labels)} inspected label value(s)." |
| if label_values else "No label values were visible in sample rows.", |
| }) |
|
|
| if {"document", "summary"}.issubset(required_fields): |
| document_fields = matched_requirements.get("document", []) |
| summary_fields = matched_requirements.get("summary", []) |
| shorter = 0 |
| for row in rows: |
| document = str(_field_value(row, document_fields) or "") |
| summary = str(_field_value(row, summary_fields) or "") |
| if document and summary and len(summary) < len(document): |
| shorter += 1 |
| tests.append({ |
| "name": "Summary shape", |
| "status": "pass" if shorter else "review", |
| "detail": "Summaries are shorter than source documents in inspected rows." |
| if shorter else "Summary/document length relationship needs review.", |
| }) |
|
|
| return tests |
|
|
|
|
| def _loader_snippet(dataset: dict[str, Any]) -> str: |
| dataset_id = dataset.get("id", "") |
| config = (dataset.get("configs") or [None])[0] |
| split = (dataset.get("splits") or ["train"])[0] |
| args = f'"{dataset_id}"' |
| if config and config != "default": |
| args += f', "{config}"' |
| return ( |
| "from datasets import load_dataset\n\n" |
| f'ds = load_dataset({args})\n' |
| f'rows = ds["{split}"]\n' |
| "print(rows[0])" |
| ) |
|
|
|
|
| def score_dataset(profile: dict[str, Any], dataset: dict[str, Any]) -> dict[str, Any]: |
| """Compute a transparent score entirely from collected evidence.""" |
| blob = _text_blob(dataset) |
| available_fields = _field_names(dataset) |
| sample_text = _sample_text(dataset) |
| evidence_blob = f"{blob} {sample_text}" |
| requested = { |
| word for word in _domain_terms(profile) |
| if word not in {"english", "labels", "label", "compact", "classifier", "dataset", "data"} |
| } |
| matched_keywords = sorted(word for word in requested if word in evidence_blob) |
| domain_check = "pass" if not requested or matched_keywords else "fail" |
|
|
| lexical_match = min(22, round(22 * len(matched_keywords) / max(2, len(requested)))) |
| task_type = profile["task_type"] |
| task_terms = _TASK_ALIASES.get(task_type, (task_type,)) |
| task_match = 13 if any(term in evidence_blob for term in task_terms) else 0 |
| relevance = lexical_match + task_match |
| modality_values = dataset.get("modalities", []) |
| modality = 15 if _contains_any(modality_values, profile["modalities"]) else 0 |
| if not modality_values: |
| modality = 7 |
| language_values = [str(value).lower() for value in dataset.get("languages", [])] |
| inferred_languages = _infer_script_languages(sample_text) if not language_values else [] |
| language_evidence = language_values or inferred_languages |
| requested_languages = { |
| value for value in profile["languages"] if value != "multilingual" |
| } |
| declared_languages = set(language_evidence) |
| if not requested_languages: |
| language = 10 |
| language_check = "pass" |
| elif not language_evidence: |
| language = 4 |
| language_check = "unknown" |
| elif requested_languages.issubset(declared_languages) or "multilingual" in declared_languages: |
| language = 10 |
| language_check = "pass" |
| elif requested_languages & declared_languages: |
| language = 5 |
| language_check = "review" |
| else: |
| language = 0 |
| language_check = "fail" |
|
|
| required_fields = profile["required_fields"] |
| proxy_label_fields = sorted(field for field in available_fields if field in _PROXY_LABEL_FIELDS) |
| embedded_label = bool( |
| required_fields |
| and ("output:" in sample_text or "intent categories" in sample_text) |
| ) |
| matched_fields = [] |
| matched_requirements = { |
| requirement: sorted(field for field in available_fields if _matches_field(requirement, field)) |
| for requirement in required_fields |
| } |
| if ( |
| {"source", "target"}.issubset(required_fields) |
| and "translation" in available_fields |
| and len(set(profile["languages"]) & available_fields) >= 2 |
| ): |
| matched_requirements["source"] = ["translation"] |
| matched_requirements["target"] = ["translation"] |
| schema_evidence = "not-required" |
| if not required_fields: |
| schema = 15 |
| elif all(matched_requirements.values()): |
| schema = 15 |
| matched_fields = sorted({ |
| field for fields_for_requirement in matched_requirements.values() |
| for field in fields_for_requirement |
| }) |
| schema_evidence = "direct" |
| elif "label" in required_fields and proxy_label_fields and all( |
| matched_requirements[requirement] for requirement in required_fields if requirement != "label" |
| ): |
| schema = 8 |
| matched_fields = proxy_label_fields |
| schema_evidence = "proxy" |
| elif "label" in required_fields and embedded_label and all( |
| matched_requirements[requirement] for requirement in required_fields if requirement != "label" |
| ): |
| schema = 5 |
| matched_fields = ["embedded instruction output"] |
| schema_evidence = "embedded" |
| else: |
| matched_count = sum(bool(fields_for_requirement) for fields_for_requirement in matched_requirements.values()) |
| schema = round(10 * matched_count / len(required_fields)) |
| matched_fields = sorted({ |
| field for fields_for_requirement in matched_requirements.values() |
| for field in fields_for_requirement |
| }) |
| schema_evidence = "missing" if available_fields else "unknown" |
|
|
| sample_tests = _sample_tests(profile, dataset, matched_requirements) |
| sample_test_passes = sum(1 for test in sample_tests if test["status"] == "pass") |
| sample_test_total = len(sample_tests) |
|
|
| license_value = dataset.get("license", "") |
| permissive = {"apache-2.0", "mit", "cc-by-4.0", "cc0-1.0", "odc-by", "bsd-3-clause"} |
| license_score = 10 if license_value in permissive else 5 if license_value else 0 |
| if profile["license"]: |
| license_check = "pass" if license_value in permissive else "unknown" if not license_value else "fail" |
| else: |
| license_check = "pass" if license_value in permissive else "unknown" if not license_value else "review" |
| documentation = 5 if dataset.get("card_complete") else 2 if dataset.get("description") else 0 |
| num_examples = int(dataset.get("num_examples") or 0) |
| if num_examples >= 10_000: |
| popularity = 5 |
| elif num_examples >= 1_000: |
| popularity = 4 |
| elif num_examples >= 100: |
| popularity = 3 |
| elif num_examples > 0: |
| popularity = 1 |
| else: |
| popularity = min( |
| 3, |
| round(math.log10(1 + dataset.get("downloads", 0) + dataset.get("likes", 0) * 20)), |
| ) |
| sample_size_adjustment = 0 |
| sample_size_check = "pass" |
| if "classification" in profile["task_type"] and num_examples: |
| if num_examples < 100: |
| sample_size_adjustment = -12 |
| sample_size_check = "review" |
| elif num_examples < 500: |
| sample_size_adjustment = -4 |
| sample_size_check = "review" |
| domain_penalty = -18 if requested and not matched_keywords else 0 |
| accessibility = 5 if dataset.get("accessible") and not dataset.get("gated") else 0 |
| total = max( |
| 0, |
| relevance + modality + language + schema + license_score |
| + documentation + popularity + accessibility + sample_size_adjustment + domain_penalty, |
| ) |
|
|
| checks = { |
| "modality": "pass" if modality == 15 else "unknown" if not modality_values else "fail", |
| "domain": domain_check, |
| "language": language_check, |
| "required_fields": "pass" if schema_evidence in {"not-required", "direct"} |
| else "review" if schema_evidence in {"proxy", "embedded"} |
| else "unknown" if schema_evidence == "unknown" else "fail", |
| "license": license_check, |
| "sample_size": sample_size_check if num_examples else "unknown", |
| "accessible": "pass" if accessibility == 5 else "fail", |
| } |
| rejection_reasons = [] |
| review_reasons = [] |
| if checks["accessible"] == "fail": |
| rejection_reasons.append("Dataset could not be inspected or is gated/private.") |
| if checks["modality"] == "fail": |
| rejection_reasons.append( |
| f"Modality {modality_values} does not match requested {profile['modalities']}." |
| ) |
| if checks["domain"] == "fail": |
| review_reasons.append( |
| "No inspected card, schema, or sample evidence matched the requested subject terms " |
| f"({', '.join(sorted(requested))})." |
| ) |
| if checks["required_fields"] == "fail": |
| review_reasons.append( |
| f"Required fields {required_fields} were not found in the inspected schema." |
| ) |
| if checks["language"] == "fail": |
| rejection_reasons.append( |
| f"Languages {language_values} do not match requested {profile['languages']}." |
| ) |
| if checks["license"] == "fail": |
| rejection_reasons.append( |
| f"License {license_value} does not meet the requested permissive/commercial constraint." |
| ) |
| recommendation_checks = ["modality", "domain", "required_fields", "accessible"] |
| if profile["languages"]: |
| recommendation_checks.append("language") |
| if profile["license"]: |
| recommendation_checks.append("license") |
| verified_core = ( |
| checks["accessible"] == "pass" |
| and checks["modality"] != "fail" |
| and checks["domain"] == "pass" |
| and checks["required_fields"] in {"pass", "review"} |
| and (not profile["languages"] or checks["language"] in {"pass", "review", "unknown"}) |
| and (not profile["license"] or checks["license"] == "pass") |
| ) |
| status = "rejected" if rejection_reasons else "recommended" if ( |
| total >= 62 |
| and verified_core |
| and schema_evidence in {"not-required", "direct"} |
| and sample_size_check == "pass" |
| ) else "conditional" |
| evidence = [ |
| f"Matched project terms: {', '.join(matched_keywords) or 'none verified'}", |
| f"Modalities: {', '.join(modality_values) or 'not declared'}", |
| f"Languages: {', '.join(language_evidence) or 'not declared'}" |
| + (" (inferred from sample script)" if inferred_languages else ""), |
| f"Features: {', '.join(sorted(available_fields)[:10]) or 'viewer schema unavailable'}", |
| f"Target evidence: {schema_evidence}" |
| + (f" ({', '.join(matched_fields)})" if matched_fields else ""), |
| f"Examples: {num_examples or 'not reported'}", |
| f"License: {license_value or 'not declared'}", |
| ] |
| strength = ( |
| f"Verified {len(matched_keywords)} project terms" |
| + (f" and fields {', '.join(matched_fields)}" if matched_fields else "") |
| + "." |
| ) |
| weakness = (rejection_reasons or review_reasons or [None])[0] or next( |
| ( |
| label for key, label in ( |
| ("domain", "The inspected metadata does not verify the requested subject domain."), |
| ("required_fields", "Required schema fields need manual confirmation."), |
| ("sample_size", "The inspected dataset is too small for reliable classifier training."), |
| ("license", "License needs manual review."), |
| ("language", "Language coverage is not declared."), |
| ) |
| if checks[key] in {"unknown", "review"} |
| ), |
| "No major metadata gap detected.", |
| ) |
| recommendation = { |
| "recommended": "Strong candidate for a first experiment; inspect sample rows before training.", |
| "conditional": "Promising candidate, but resolve the highlighted evidence gaps first.", |
| "rejected": "Do not use for this request unless the project requirements change.", |
| }[status] |
| quality = round( |
| (schema + license_score + documentation + popularity + accessibility) / 40 * 100 |
| ) |
| low_adoption = (dataset.get("downloads", 0) < 2_000 and dataset.get("likes", 0) < 25) |
| hidden_gem = ( |
| status != "rejected" |
| and total >= 60 |
| and low_adoption |
| and checks["domain"] == "pass" |
| and checks["required_fields"] in {"pass", "review"} |
| and checks["accessible"] == "pass" |
| ) |
| badges = ["hidden_gem"] if hidden_gem else [] |
| discovery_note = ( |
| "Hidden gem: low adoption, but the inspected evidence fits this brief." |
| if hidden_gem else "" |
| ) |
| return { |
| **dataset, |
| "score": total, |
| "relevance": round((relevance + modality + language) / 60 * 100), |
| "quality": min(100, quality), |
| "status": status, |
| "score_breakdown": { |
| "project_match": relevance, |
| "modality": modality, |
| "language": language, |
| "schema": schema, |
| "license": license_score, |
| "documentation": documentation, |
| "adoption": popularity, |
| "sample_size_adjustment": sample_size_adjustment, |
| "domain_penalty": domain_penalty, |
| "accessibility": accessibility, |
| }, |
| "checks": checks, |
| "sample_tests": sample_tests, |
| "sample_test_summary": f"{sample_test_passes}/{sample_test_total} sample tests passed", |
| "schema_evidence": schema_evidence, |
| "evidence": evidence, |
| "rejection_reasons": rejection_reasons, |
| "review_reasons": review_reasons, |
| "strength": strength, |
| "weakness": weakness, |
| "recommendation": recommendation, |
| "badges": badges, |
| "discovery_note": discovery_note, |
| "loader_snippet": _loader_snippet(dataset), |
| } |
|
|
|
|
| def _cross_reference(datasets: list[dict[str, Any]]) -> list[dict[str, str]]: |
| candidates = [dataset for dataset in datasets if dataset["status"] != "rejected"][:6] |
| pairs = [] |
| for index, first in enumerate(candidates): |
| for second in candidates[index + 1:]: |
| first_terms = set(first.get("modalities", []) + first.get("languages", [])) |
| second_terms = set(second.get("modalities", []) + second.get("languages", [])) |
| if first_terms != second_terms or set(first.get("features", [])) != set(second.get("features", [])): |
| pairs.append({ |
| "from": first["id"], |
| "to": second["id"], |
| "label": "complementary coverage", |
| }) |
| if len(pairs) >= 6: |
| return pairs |
| return pairs |
|
|
|
|
| def _rank_key(profile: dict[str, Any], dataset: dict[str, Any]) -> tuple[int, int, int, int]: |
| checks = dataset["checks"] |
| hidden_gem = 1 if "hidden_gem" in dataset.get("badges", []) else 0 |
| sample_passes = sum(1 for test in dataset.get("sample_tests", []) if test.get("status") == "pass") |
| evidence_fit = ( |
| (2 if checks["required_fields"] == "pass" else 0) |
| + (1 if checks["domain"] == "pass" else 0) |
| + (2 if profile["languages"] and checks["language"] == "pass" else 0) |
| + (2 if profile["license"] and checks["license"] == "pass" else 0) |
| + sample_passes |
| ) |
| status_rank = {"recommended": 2, "conditional": 1, "rejected": 0}[dataset["status"]] |
| return status_rank, hidden_gem, evidence_fit, dataset["score"] |
|
|
|
|
| def _select_candidates_for_profile( |
| profile: dict[str, Any], |
| collected: dict[str, dict[str, Any]], |
| search_batches: list[list[str]], |
| limit: int, |
| ) -> list[dict[str, Any]]: |
| global_ranked = sorted( |
| collected.values(), |
| key=lambda dataset: _pre_score(profile, dataset), |
| reverse=True, |
| ) |
| diversified_ids: list[str] = [] |
| diversified_seen: set[str] = set() |
| for position in range(8): |
| for batch in search_batches: |
| if position >= len(batch): |
| continue |
| dataset_id = batch[position] |
| if dataset_id not in diversified_seen: |
| diversified_seen.add(dataset_id) |
| diversified_ids.append(dataset_id) |
| for dataset in global_ranked: |
| if dataset["id"] not in diversified_seen: |
| diversified_ids.append(dataset["id"]) |
| return [ |
| collected[dataset_id] |
| for dataset_id in diversified_ids[:limit] |
| if dataset_id in collected |
| ] |
|
|
|
|
| def _reflect_on_results( |
| profile: dict[str, Any], |
| tried_queries: list[str], |
| inspected: list[dict[str, Any]], |
| ) -> dict[str, Any]: |
| if not inspected: |
| summary = "The first search did not produce inspectable datasets, so the agent is broadening the query language." |
| strategy = "broaden search" |
| else: |
| failures = { |
| key: sum(1 for dataset in inspected if dataset.get("checks", {}).get(key) in {"fail", "unknown"}) |
| for key in ("domain", "required_fields", "language", "license", "accessible") |
| } |
| worst = max(failures, key=failures.get) |
| if worst == "required_fields": |
| summary = "The first pass found topical datasets, but too many missed the requested schema." |
| strategy = "schema-first search" |
| elif worst == "domain": |
| summary = "The first pass found task-shaped datasets, but several were off-topic." |
| strategy = "domain-first search" |
| elif worst == "language": |
| summary = "The first pass found candidates with weak language evidence." |
| strategy = "language-focused search" |
| else: |
| summary = "The first pass found partial fits; the agent is looking for less obvious alternatives." |
| strategy = "hidden-gem search" |
|
|
| terms = _domain_terms(profile) |
| domain = terms[:2] |
| required = profile["required_fields"] |
| next_queries: list[str] = [] |
| if domain and {"question", "answer"}.issubset(required): |
| next_queries.extend([ |
| f"{domain[0]} qa dataset", |
| " ".join(domain + ["question answer"]), |
| " ".join(domain + ["benchmark"]), |
| ]) |
| if domain and "label" in required: |
| next_queries.extend([ |
| f"{domain[0]} labeled dataset", |
| f"{domain[0]} intent labels", |
| ]) |
| if domain and "summary" in required: |
| next_queries.extend([ |
| " ".join(domain + ["summaries"]), |
| " ".join(domain + ["summarization dataset"]), |
| ]) |
| if profile["languages"] and domain: |
| next_queries.append(" ".join([profile["languages"][0], domain[0], "dataset"])) |
| if domain: |
| next_queries.append(" ".join(domain + ["data"])) |
| next_queries.extend(generate_queries("", profile)) |
| cleaned = [] |
| for query in next_queries: |
| normalized = re.sub(r"\s+", " ", query).strip() |
| if normalized and normalized not in tried_queries and normalized not in cleaned: |
| cleaned.append(normalized) |
| return { |
| "summary": summary, |
| "strategy": strategy, |
| "next_queries": cleaned[:4], |
| "reason": "Reflection is based on failed checks from the first inspected batch.", |
| } |
|
|
|
|
| def weave_events(task: str, max_datasets: int = 8) -> Iterator[dict[str, Any]]: |
| task = task.strip() |
| if not task: |
| raise ValueError("Task description is required.") |
| if len(task) > MAX_TASK_LENGTH: |
| raise ValueError(f"Task description must be {MAX_TASK_LENGTH} characters or fewer.") |
|
|
| started = time.time() |
| yield {"type": "started", "task": task, "message": "Research session started."} |
| profile, llm_used = parse_task(task) |
| queries = generate_queries(task, profile) |
| yield { |
| "type": "plan", |
| "profile": profile, |
| "queries": queries, |
| "model_used": MODEL if llm_used else None, |
| "fallback_used": not llm_used, |
| "message": "Converted the request into explicit dataset requirements.", |
| } |
|
|
| collected: dict[str, dict[str, Any]] = {} |
| search_batches: list[list[str]] = [] |
| inspected: list[dict[str, Any]] = [] |
| inspected_ids: set[str] = set() |
|
|
| first_pass_queries = queries[:5] |
| reserve_queries = queries[5:] |
|
|
| for query in first_pass_queries: |
| found = search_datasets(query, limit=35) |
| search_batches.append([dataset["id"] for dataset in found]) |
| for dataset in found: |
| current = collected.get(dataset["id"]) |
| if current is None or _pre_score(profile, dataset) > _pre_score(profile, current): |
| collected[dataset["id"]] = dataset |
| yield { |
| "type": "search", |
| "query": query, |
| "found": len(found), |
| "unique": len(collected), |
| "message": f"Searched “{query}” and found {len(found)} candidates.", |
| } |
|
|
| inspection_limit = max(max_datasets * 4, 28) |
| first_pass_limit = max(max_datasets * 2, 16) |
| pre_ranked = _select_candidates_for_profile( |
| profile, |
| collected, |
| search_batches, |
| first_pass_limit, |
| ) |
| yield { |
| "type": "search", |
| "query": "first evidence pool", |
| "found": len(pre_ranked), |
| "unique": len(collected), |
| "message": f"Prepared {len(pre_ranked)} diverse candidates for first-pass inspection.", |
| } |
| with ThreadPoolExecutor(max_workers=min(4, max(1, len(pre_ranked)))) as pool: |
| futures = { |
| pool.submit(inspect_dataset, dataset["id"], dataset): dataset["id"] |
| for dataset in pre_ranked |
| } |
| for future in as_completed(futures): |
| dataset_id = futures[future] |
| try: |
| evidence = future.result() |
| except Exception as exc: |
| evidence = { |
| **next(item for item in pre_ranked if item["id"] == dataset_id), |
| "accessible": False, |
| "inspection_error": str(exc), |
| "features": [], |
| "sample_rows": [], |
| "configs": [], |
| "splits": [], |
| } |
| scored = score_dataset(profile, evidence) |
| inspected.append(scored) |
| inspected_ids.add(dataset_id) |
| yield { |
| "type": "inspect", |
| "dataset_id": dataset_id, |
| "status": scored["status"], |
| "score": scored["score"], |
| "checks": scored["checks"], |
| "message": f"Inspected {dataset_id}: {scored['status']} ({scored['score']}/100).", |
| } |
| yield {"type": "candidate", "dataset": _public_dataset(scored)} |
|
|
| reflection = _reflect_on_results(profile, first_pass_queries, inspected) |
| second_pass_queries = list(dict.fromkeys(reflection["next_queries"] + reserve_queries))[:4] |
| yield { |
| "type": "reflect", |
| "summary": reflection["summary"], |
| "strategy": reflection["strategy"], |
| "next_queries": second_pass_queries, |
| "message": f"{reflection['summary']} Trying {len(second_pass_queries)} deeper search angle(s).", |
| } |
|
|
| for query in second_pass_queries: |
| found = search_datasets(query, limit=35) |
| search_batches.append([dataset["id"] for dataset in found]) |
| new_count = 0 |
| for dataset in found: |
| current = collected.get(dataset["id"]) |
| if current is None: |
| new_count += 1 |
| if current is None or _pre_score(profile, dataset) > _pre_score(profile, current): |
| collected[dataset["id"]] = dataset |
| yield { |
| "type": "search", |
| "query": query, |
| "found": len(found), |
| "unique": len(collected), |
| "message": f"Deepened with “{query}” and found {len(found)} candidates ({new_count} new).", |
| } |
|
|
| deeper_pool = [ |
| dataset for dataset in _select_candidates_for_profile( |
| profile, |
| collected, |
| search_batches, |
| inspection_limit, |
| ) |
| if dataset["id"] not in inspected_ids |
| ][:max(max_datasets * 2, 16)] |
| yield { |
| "type": "search", |
| "query": "deep candidate pool", |
| "found": len(deeper_pool), |
| "unique": len(collected), |
| "message": f"Prepared {len(deeper_pool)} fresh candidates after reflection.", |
| } |
| if deeper_pool: |
| with ThreadPoolExecutor(max_workers=min(4, len(deeper_pool))) as pool: |
| futures = { |
| pool.submit(inspect_dataset, dataset["id"], dataset): dataset["id"] |
| for dataset in deeper_pool |
| } |
| for future in as_completed(futures): |
| dataset_id = futures[future] |
| try: |
| evidence = future.result() |
| except Exception as exc: |
| evidence = { |
| **next(item for item in deeper_pool if item["id"] == dataset_id), |
| "accessible": False, |
| "inspection_error": str(exc), |
| "features": [], |
| "sample_rows": [], |
| "configs": [], |
| "splits": [], |
| } |
| scored = score_dataset(profile, evidence) |
| inspected.append(scored) |
| inspected_ids.add(dataset_id) |
| yield { |
| "type": "inspect", |
| "dataset_id": dataset_id, |
| "status": scored["status"], |
| "score": scored["score"], |
| "checks": scored["checks"], |
| "message": f"Inspected {dataset_id}: {scored['status']} ({scored['score']}/100).", |
| } |
| yield {"type": "candidate", "dataset": _public_dataset(scored)} |
|
|
| ranked = sorted( |
| inspected, |
| key=lambda dataset: _rank_key(profile, dataset), |
| reverse=True, |
| ) |
| ranked = ranked[:max_datasets] |
| pairs = _cross_reference(ranked) |
| nodes = [ |
| { |
| "id": dataset["id"], |
| "score": dataset["score"], |
| "status": dataset["status"], |
| "downloads": dataset.get("downloads", 0), |
| } |
| for dataset in ranked |
| ] |
| result = { |
| "task": task, |
| "profile": profile, |
| "queries": queries, |
| "datasets": [_public_dataset(dataset) for dataset in ranked], |
| "nodes": nodes, |
| "threads": pairs, |
| "top_pick": next( |
| (dataset["id"] for dataset in ranked if dataset["status"] != "rejected"), |
| ranked[0]["id"] if ranked else None, |
| ), |
| "model_used": MODEL if llm_used else None, |
| "fallback_used": not llm_used, |
| "elapsed_ms": round((time.time() - started) * 1000), |
| "reflection": reflection, |
| } |
| yield { |
| "type": "ranking", |
| "top_pick": result["top_pick"], |
| "count": len(ranked), |
| "message": f"Ranked {len(ranked)} inspected candidates using verified evidence.", |
| } |
| yield {"type": "complete", "result": result, "message": "Dataset research complete."} |
|
|
|
|
| def _public_dataset(dataset: dict[str, Any]) -> dict[str, Any]: |
| allowed = { |
| "id", "author", "description", "downloads", "likes", "tags", |
| "task_categories", "languages", "license", "size_category", "formats", |
| "modalities", "configs", "splits", "features", "sample_rows", "hub_url", |
| "accessible", "inspection_error", "card_complete", "num_examples", "score", "relevance", |
| "quality", "status", "score_breakdown", "checks", "evidence", |
| "schema_evidence", "rejection_reasons", "review_reasons", "strength", "weakness", "recommendation", |
| "sample_tests", "sample_test_summary", "badges", "discovery_note", "loader_snippet", |
| } |
| return {key: value for key, value in dataset.items() if key in allowed} |
|
|
|
|
| def weave(task: str, max_datasets: int = 8) -> dict[str, Any]: |
| result = None |
| for event in weave_events(task, max_datasets=max_datasets): |
| if event["type"] == "complete": |
| result = event["result"] |
| return result or { |
| "task": task, |
| "datasets": [], |
| "nodes": [], |
| "threads": [], |
| "top_pick": None, |
| "fallback_used": True, |
| } |
|
|
