wehub-resource-sync
113 lines
4.6 KiB
Python
113 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
from skillopt.datasets.base import BatchSpec
|
|
from skillopt.envs.base import EnvAdapter
|
|
from skillopt.envs.officeqa.dataloader import OfficeQADataLoader
|
|
from skillopt.envs.officeqa.rollout import run_batch
|
|
|
|
|
|
class OfficeQAAdapter(EnvAdapter):
|
|
def __init__(
|
|
self,
|
|
split_dir: str = "",
|
|
data_path: str = "",
|
|
split_mode: str = "split_dir",
|
|
split_ratio: str = "2:1:7",
|
|
split_seed: int = 42,
|
|
split_output_dir: str = "",
|
|
workers: int = 8,
|
|
analyst_workers: int = 8,
|
|
failure_only: bool = False,
|
|
minibatch_size: int = 8,
|
|
edit_budget: int = 4,
|
|
seed: int = 42,
|
|
limit: int = 0,
|
|
max_tool_turns: int = 12,
|
|
max_completion_tokens: int = 16384,
|
|
search_mode: str = "offline",
|
|
max_queries_per_turn: int = 4,
|
|
search_api_url: str = os.environ.get("OFFICEQA_SEARCH_API_URL", "http://localhost:8080/search_tool/search"),
|
|
search_auth_env: str = "OFFICEQA_CUSTOM_SEARCH_AUTH",
|
|
search_provider: str = "duckduckgo",
|
|
search_max_num_results: int = 4,
|
|
search_timeout_seconds: int = 20,
|
|
use_local_tools: bool = True,
|
|
data_dirs: list[str] | str | None = None,
|
|
docs_dirs: list[str] | str | None = None, ) -> None:
|
|
self.workers = workers
|
|
self.analyst_workers = analyst_workers
|
|
self.failure_only = failure_only
|
|
self.minibatch_size = minibatch_size
|
|
self.edit_budget = edit_budget
|
|
self.max_tool_turns = max_tool_turns
|
|
self.max_completion_tokens = int(max_completion_tokens)
|
|
self.search_mode = str(search_mode or "offline")
|
|
self.max_queries_per_turn = int(max_queries_per_turn)
|
|
self.search_api_url = str(search_api_url or "").strip()
|
|
self.search_auth_env = str(search_auth_env or "OFFICEQA_CUSTOM_SEARCH_AUTH").strip()
|
|
self.search_provider = str(search_provider or "duckduckgo").strip()
|
|
self.search_max_num_results = int(search_max_num_results)
|
|
self.search_timeout_seconds = int(search_timeout_seconds)
|
|
self.use_local_tools = bool(use_local_tools)
|
|
self.data_dirs = data_dirs if data_dirs is not None else docs_dirs
|
|
self.dataloader = OfficeQADataLoader(
|
|
split_dir=split_dir,
|
|
data_path=data_path,
|
|
split_mode=split_mode,
|
|
split_ratio=split_ratio,
|
|
split_seed=split_seed,
|
|
split_output_dir=split_output_dir,
|
|
seed=seed,
|
|
limit=limit,
|
|
)
|
|
|
|
def setup(self, cfg: dict) -> None:
|
|
super().setup(cfg)
|
|
self.dataloader.setup(cfg)
|
|
|
|
def get_dataloader(self):
|
|
return self.dataloader
|
|
|
|
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
|
return list(batch.payload or [])
|
|
|
|
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
|
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
|
return self.build_env_from_batch(batch, **kwargs)
|
|
|
|
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
|
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
|
return self.build_env_from_batch(batch, **kwargs)
|
|
|
|
def rollout(self, env_manager, skill_content: str, out_dir: str, **kwargs) -> list[dict]:
|
|
items: list[dict] = env_manager
|
|
return run_batch(
|
|
items=items,
|
|
out_root=out_dir,
|
|
skill_content=skill_content,
|
|
workers=self.workers,
|
|
max_tool_turns=self.max_tool_turns,
|
|
max_completion_tokens=self.max_completion_tokens,
|
|
search_mode=self.search_mode,
|
|
max_queries_per_turn=self.max_queries_per_turn,
|
|
search_api_url=self.search_api_url,
|
|
search_auth_env=self.search_auth_env,
|
|
search_provider=self.search_provider,
|
|
search_max_num_results=self.search_max_num_results,
|
|
search_timeout_seconds=self.search_timeout_seconds,
|
|
use_local_tools=self.use_local_tools,
|
|
data_dirs=self.data_dirs,
|
|
diagnostic_mode=kwargs.get("diagnostic_mode", False),
|
|
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
|
|
)
|
|
|
|
def get_task_types(self) -> list[str]:
|
|
seen: list[str] = []
|
|
for item in self.dataloader.train_items + self.dataloader.val_items + self.dataloader.test_items:
|
|
task_type = str(item.get("task_type") or "officeqa")
|
|
if task_type not in seen:
|
|
seen.append(task_type)
|
|
return seen or ["officeqa"]
|