File size: 3,643 Bytes
b0fbec3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """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
|