| import torch |
| import random |
| from torch.utils.data import IterableDataset, DataLoader |
| from datasets import load_dataset |
| from transformers import PreTrainedTokenizerFast |
| from typing import Optional, Iterator |
|
|
| LANGUAGES = { |
| "python": 0.35, "javascript": 0.20, "typescript": 0.15, |
| "cpp": 0.10, "rust": 0.08, "go": 0.07, "java": 0.05, |
| } |
| LANG_TOKEN_MAP = { |
| "python": "<|python|>", "javascript": "<|javascript|>", |
| "typescript": "<|typescript|>", "cpp": "<|cpp|>", |
| "rust": "<|rust|>", "go": "<|go|>", "java": "<|java|>", |
| } |
|
|
| class TheStackStreamDataset(IterableDataset): |
| def __init__(self, tokenizer, max_length=2048, languages=None, |
| split="train", max_samples_per_lang=500_000, fim_rate=0.5): |
| self.tokenizer = tokenizer |
| self.max_length = max_length |
| self.languages = languages or list(LANGUAGES.keys()) |
| self.split = split |
| self.max_samples = max_samples_per_lang |
| self.fim_rate = fim_rate |
|
|
| def _get_lang_dataset(self, lang): |
| return load_dataset( |
| "bigcode/the-stack", data_dir=f"data/{lang}", |
| split=self.split, streaming=True, trust_remote_code=True, |
| ) |
|
|
| def _tokenize(self, code, lang): |
| text = f"{LANG_TOKEN_MAP.get(lang, '')}{code}" |
| tokens = self.tokenizer(text, max_length=self.max_length, truncation=True) |
| ids = tokens["input_ids"] |
| if len(ids) < 64: |
| return None |
| return {"input_ids": torch.tensor(ids, dtype=torch.long), |
| "labels": torch.tensor(ids, dtype=torch.long)} |
|
|
| def _apply_fim(self, code): |
| if random.random() > self.fim_rate: |
| return code |
| lines = code.split("\n") |
| if len(lines) < 4: |
| return code |
| start = random.randint(1, len(lines) - 3) |
| end = random.randint(start + 1, min(start + 10, len(lines) - 1)) |
| prefix = "\n".join(lines[:start]) |
| middle = "\n".join(lines[start:end]) |
| suffix = "\n".join(lines[end:]) |
| return f"<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>{middle}" |
|
|
| def __iter__(self): |
| datasets = {} |
| for lang in self.languages: |
| try: |
| datasets[lang] = iter(self._get_lang_dataset(lang)) |
| except Exception as e: |
| print(f"Warning: could not load {lang}: {e}") |
| lang_list = list(datasets.keys()) |
| weights = [LANGUAGES.get(l, 1.0) for l in lang_list] |
| counts = {l: 0 for l in lang_list} |
| while lang_list: |
| lang = random.choices(lang_list, weights=weights, k=1)[0] |
| try: |
| sample = next(datasets[lang]) |
| code = sample.get("content", "") |
| if not code.strip(): |
| continue |
| code = self._apply_fim(code) |
| item = self._tokenize(code, lang) |
| if item: |
| counts[lang] += 1 |
| yield item |
| if counts[lang] >= self.max_samples: |
| lang_list.remove(lang) |
| weights = [LANGUAGES.get(l, 1.0) for l in lang_list] |
| except StopIteration: |
| lang_list.remove(lang) |
| weights = [LANGUAGES.get(l, 1.0) for l in lang_list] |
|
|
| class CodeCollator: |
| def __init__(self, pad_token_id=0, max_length=2048): |
| self.pad_id = pad_token_id |
| self.max_length = max_length |
|
|
| def __call__(self, batch): |
| max_len = min(max(len(x["input_ids"]) for x in batch), self.max_length) |
| input_ids = torch.zeros(len(batch), max_len, dtype=torch.long) |
| labels = torch.full((len(batch), max_len), -100, dtype=torch.long) |
| attention_mask = torch.zeros(len(batch), max_len, dtype=torch.long) |
| for i, item in enumerate(batch): |
| length = min(len(item["input_ids"]), max_len) |
| input_ids[i, :length] = item["input_ids"][:length] |
| labels[i, :length] = item["labels"][:length] |
| attention_mask[i, :length] = 1 |
| return {"input_ids": input_ids, "labels": labels, "attention_mask": attention_mask} |
