41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
import json
|
|
|
|
from torch.utils.data import Dataset
|
|
|
|
from agentic.data_analysis_env import AGENTIC_SYSTEM_PROMPT, format_agentic_user_prompt, get_agentic_tools
|
|
|
|
|
|
def load_agentic_jsonl(path):
|
|
rows = []
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
rows.append(json.loads(line))
|
|
return rows
|
|
|
|
|
|
class AgenticRLDataset(Dataset):
|
|
def __init__(self, jsonl_path, tokenizer=None, max_length=2048):
|
|
super().__init__()
|
|
self.samples = load_agentic_jsonl(jsonl_path)
|
|
self.tokenizer = tokenizer
|
|
self.max_length = max_length
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, index):
|
|
sample = self.samples[index]
|
|
tools = get_agentic_tools(sample.get("tools"))
|
|
messages = [
|
|
{"role": "system", "content": AGENTIC_SYSTEM_PROMPT},
|
|
{"role": "user", "content": format_agentic_user_prompt(sample)},
|
|
]
|
|
return {
|
|
"id": sample.get("id", str(index)),
|
|
"messages": messages,
|
|
"tools": tools,
|
|
"sample": sample,
|
|
}
|