File size: 1,137 Bytes
c6bc767 | 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 | import json
import torch
from torch.utils.data import Dataset
class AgentRLDataset(Dataset):
def __init__(self, jsonl_path, tokenizer, max_length=1024):
super().__init__()
self.tokenizer = tokenizer
self.max_length = max_length
self.samples = []
with open(jsonl_path, 'r', encoding='utf-8') as f:
for line in f:
self.samples.append(json.loads(line.strip()))
def __len__(self):
return len(self.samples)
def parse_conversations(self, conversations):
messages = []
tools = None
for message in conversations:
message = dict(message)
if message.get("role") == "system" and message.get("tools"):
tools = json.loads(message["tools"]) if isinstance(message["tools"], str) else message["tools"]
messages.append(message)
return messages[:-1], tools
def __getitem__(self, index):
sample = self.samples[index]
messages, tools = self.parse_conversations(sample['conversations'])
return {'messages': messages, 'tools': tools, 'gt': sample['gt']}
|