chimp/src/dataset.py

class CustomDataset:
 def __init__(self, data):
 self.features = ['instruction', 'context', 'response']
 self.num_rows = len(data)
 self.data = data

 def __getitem__(self, idx):
 if idx < 0 or idx >= self.num_rows:
 raise IndexError("Index out of range")
 return {
 'instruction': self.data[idx]['instruction'],
 'context': self.data[idx]['context'],
 'response': self.data[idx]['response']
 }

 def __repr__(self):
 return f"Dataset({{'features': {self.features}, 'num_rows': {self.num_rows}}})"


def format_data(sample):
 instruction = f"<s>[INST] {sample['instruction']}"
 context = f"Here's some context: {sample['context']}" if len(sample["context"]) > 0 else None
 response = f" [/INST] {sample['response']}"
 # join all the parts together
 prompt = "".join([i for i in [instruction, context, response] if i is not None])
 return prompt

# template dataset to add prompt to each sample
def template_dataset(sample, tokenizer):
 sample["text"] = f"{format_data(sample)}{tokenizer.eos_token}"
 return sample