File size: 3,050 Bytes
166edf7 d1ea2f2 166edf7 d1ea2f2 166edf7 f550456 166edf7 0c2ae95 166edf7 d1ea2f2 166edf7 d1ea2f2 8b94669 d1ea2f2 166edf7 0c2ae95 166edf7 0c2ae95 d1ea2f2 166edf7 | 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 | import os
import json
import torch
import random
from torch.utils.data import Dataset
from transformers import BertTokenizer
class MixedDataset(Dataset):
def __init__(self, bert_model, stage, anno_file, tool_capacity, seed):
self.stage = stage
self.seed = seed
self.tool_capacity = tool_capacity
self.tools, self.samples = self.load_data(anno_file)
self.tool_ids = list(self.tools.keys())
self.tokenizer = BertTokenizer.from_pretrained(bert_model)
def load_data(self, anno_file):
with open(anno_file, "r") as f:
data = json.load(f)
tools = data["tools"]
samples = data["samples"]
tools = {tool["id"]: tool for tool in tools}
return tools, samples
def encode_text(self, text, padding=True):
if padding:
inputs = self.tokenizer(
text,
max_length=128,
padding="max_length",
truncation=True,
)
else:
inputs = self.tokenizer(
text,
max_length=128,
truncation=True,
)
ids = torch.tensor(inputs["input_ids"], dtype=torch.long)
mask = torch.tensor(inputs["attention_mask"], dtype=torch.long)
return ids, mask
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
sample = self.samples[idx]
inst = sample["instruction"]
inst_ids, inst_mask = self.encode_text(inst)
if self.stage == "train":
tool_id = random.choice(sample["tools"])
tool_desc = self.tools[tool_id]["description"]
tool_desc_ids, tool_desc_mask = self.encode_text(tool_desc)
return {
"inst_ids": inst_ids,
"inst_mask": inst_mask,
"tool_ids": tool_desc_ids,
"tool_mask": tool_desc_mask,
}
else:
# for testing, we sample a random set of tools + the correct tool, size = tool_capacity
# wrong tools are sampled randomly from self.tools
correct_tools = sample["tools"]
random.seed(self.seed + idx)
wrong_tools = random.sample(
[tool for tool in self.tool_ids if tool not in correct_tools],
self.tool_capacity - len(correct_tools),
)
correct_tool_mask = torch.tensor(
[1] * len(correct_tools)
+ [0] * (self.tool_capacity - len(correct_tools)),
dtype=torch.bool,
)
tools = correct_tools + wrong_tools
tool_ids, tool_mask = self.encode_text(
[self.tools[tool_id]["description"] for tool_id in tools]
)
return {
"inst_ids": inst_ids,
"inst_mask": inst_mask,
"tool_ids": tool_ids,
"tool_mask": tool_mask,
"correct_tool_mask": correct_tool_mask,
}
|