| """ |
| PACE Actor Dataset |
| Task: Ingests clean Python code, tokenizes it, and duplicates the input_ids to act as causal LM labels. |
| """ |
|
|
| import pandas as pd |
| from torch.utils.data import Dataset, DataLoader |
| from transformers import AutoTokenizer |
|
|
|
|
| class CoderActorDataset(Dataset): |
| def __init__(self, parquet_path, tokenizer, max_length=512): |
| |
| |
| |
| self.df = pd.read_parquet(parquet_path) |
| self.tokenizer = tokenizer |
| self.max_length = max_length |
|
|
| def __len__(self): |
| return len(self.df) |
|
|
| def __getitem__(self, idx): |
| |
| |
| |
| clean_code_text = self.df.iloc[idx]["content"] |
|
|
| |
| |
| |
| inputs = self.tokenizer( |
| clean_code_text, |
| max_length=self.max_length, |
| truncation=True, |
| padding="max_length", |
| return_tensors="pt", |
| ) |
|
|
| input_ids = inputs["input_ids"].squeeze(0) |
| attention_mask = inputs["attention_mask"].squeeze(0) |
|
|
| |
| |
| |
| return { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| "labels": input_ids.clone(), |
| } |
|
|
|
|
| |
| |
| |
| if __name__ == "__main__": |
| tokenizer = AutoTokenizer.from_pretrained("bigcode/tiny_starcoder_py") |
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token |
|
|
| test_path = "masteries/coding/data/raw/stack_python_157k.parquet" |
| dataset = CoderActorDataset(test_path, tokenizer, max_length=32) |
| print(f"Total rows: {len(dataset)}") |
|
|
| dataloader = DataLoader(dataset, batch_size=4, shuffle=True) |
| for batch in dataloader: |
| print("Input IDs Shape:", batch["input_ids"].shape) |
| print("Labels Shape:", batch["labels"].shape) |
| break |
|
|
