| """Task sampler: loads the seed corpus and samples Tasks by difficulty. |
| |
| Difficulty is auto-derived from script line count. Category is auto-detected |
| from script content (text_classification, ner, translation, etc.). |
| """ |
| from __future__ import annotations |
|
|
| import random |
| from pathlib import Path |
| from typing import Optional |
|
|
| from forgeenv.tasks.models import Task |
|
|
|
|
| def _detect_category(content: str) -> str: |
| cl = content.lower() |
| if "sequenceclassification" in cl or "sentiment" in cl or "ag_news" in cl or "sst2" in cl: |
| return "text_classification" |
| if "tokenclassification" in cl or "ner" in cl or "conll" in cl: |
| return "ner" |
| if "seq2seq" in cl or "translation" in cl or "summariz" in cl or "t5" in cl: |
| return "seq2seq" |
| if "causallm" in cl or "gpt2" in cl or "wikitext" in cl: |
| return "text_generation" |
| if "imageclassification" in cl or "vit" in cl or "cifar" in cl or "mnist" in cl: |
| return "image_classification" |
| if "questionanswering" in cl or "squad" in cl: |
| return "qa" |
| if "logisticregression" in cl or "make_classification" in cl: |
| return "tabular" |
| if "regression" in cl: |
| return "regression" |
| return "general" |
|
|
|
|
| def _derive_difficulty(content: str) -> str: |
| lines = len(content.splitlines()) |
| if lines < 30: |
| return "easy" |
| if lines < 60: |
| return "medium" |
| return "hard" |
|
|
|
|
| class TaskSampler: |
| """Loads seed corpus and samples tasks by difficulty / category.""" |
|
|
| def __init__(self, seed_dir: Optional[str] = None) -> None: |
| if seed_dir is None: |
| seed_dir = str(Path(__file__).parent / "seed_corpus") |
|
|
| self.tasks: list[Task] = [] |
| self._load_corpus(seed_dir) |
|
|
| def _load_corpus(self, seed_dir: str) -> None: |
| corpus_path = Path(seed_dir) |
| if not corpus_path.exists(): |
| return |
|
|
| for py_file in sorted(corpus_path.glob("*.py")): |
| if py_file.name.startswith("__"): |
| continue |
|
|
| content = py_file.read_text(encoding="utf-8") |
| task_id = py_file.stem |
| difficulty = _derive_difficulty(content) |
| category = _detect_category(content) |
|
|
| description = "" |
| if content.startswith('"""'): |
| end = content.find('"""', 3) |
| if end != -1: |
| description = content[3:end].strip() |
|
|
| self.tasks.append( |
| Task( |
| task_id=task_id, |
| description=description or f"Training script: {task_id}", |
| script_content=content, |
| difficulty=difficulty, |
| category=category, |
| ) |
| ) |
|
|
| def sample(self, difficulty: Optional[str] = None) -> Optional[Task]: |
| candidates = self.tasks |
| if difficulty is not None: |
| filtered = [t for t in self.tasks if t.difficulty == difficulty] |
| if filtered: |
| candidates = filtered |
| return random.choice(candidates) if candidates else None |
|
|
| def sample_batch( |
| self, n: int, difficulty: Optional[str] = None |
| ) -> list[Task]: |
| return [t for t in (self.sample(difficulty) for _ in range(n)) if t is not None] |
|
|
| def get_all_categories(self) -> list[str]: |
| return sorted({t.category for t in self.tasks}) |
|
|
| def get_by_id(self, task_id: str) -> Optional[Task]: |
| for t in self.tasks: |
| if t.task_id == task_id: |
| return t |
| return None |
|
|
