Create data/dataset.py

data/dataset.py

@@ -0,0 +1,100 @@
+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}