| |
| import json, os, sys, torch |
| from pathlib import Path |
| from datasets import load_dataset |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from train_ultra import CharTokenizer |
|
|
| DATA_DIR = "/root/cognet-1b/data_1b" |
| os.makedirs(DATA_DIR, exist_ok=True) |
|
|
| tokenizer = CharTokenizer.load("/root/cognet-1b/tokenizer_v3.json") |
|
|
| def tokenize_texts(texts, desc=""): |
| all_ids = [] |
| for i, text in enumerate(texts): |
| if not text or len(text.strip()) < 10: |
| continue |
| all_ids.extend(tokenizer.encode(text)) |
| if i % 50000 == 0 and i > 0: |
| print(f" {desc}: {i:,} texts -> {len(all_ids):,} tokens") |
| return all_ids |
|
|
| all_token_ids = [] |
|
|
| |
| print("1/5 - WikiText-103-RAW...") |
| try: |
| ds = load_dataset("wikitext", "wikitext-103-raw-v1", split="train", trust_remote_code=True) |
| texts = [x["text"] for x in ds if x["text"].strip()] |
| ids = tokenize_texts(texts, "WikiText-103") |
| all_token_ids.extend(ids) |
| print(f" OK WikiText-103: {len(ids):,} tokens") |
| del ds, texts |
| except Exception as e: |
| print(f" FAIL WikiText-103: {e}") |
|
|
| |
| print("2/5 - Alpaca...") |
| try: |
| ds = load_dataset("tatsu-lab/alpaca", split="train", trust_remote_code=True) |
| texts = [] |
| for x in ds: |
| if x["input"]: |
| t = f"### Instruction:\n{x['instruction']}\n\n### Input:\n{x['input']}\n\n### Response:\n{x['output']}" |
| else: |
| t = f"### Instruction:\n{x['instruction']}\n\n### Response:\n{x['output']}" |
| texts.append(t) |
| ids = tokenize_texts(texts, "Alpaca") |
| all_token_ids.extend(ids) |
| print(f" OK Alpaca: {len(ids):,} tokens") |
| del ds, texts |
| except Exception as e: |
| print(f" FAIL Alpaca: {e}") |
|
|
| |
| print("3/5 - OSCAR French...") |
| try: |
| ds = load_dataset("oscar", "unshuffled_deduplicated_fr", split="train", streaming=True, trust_remote_code=True) |
| texts = [] |
| count = 0 |
| for x in ds: |
| texts.append(x["text"]) |
| count += 1 |
| if count >= 200000: |
| break |
| ids = tokenize_texts(texts, "OSCAR-FR") |
| all_token_ids.extend(ids) |
| print(f" OK OSCAR-FR: {len(ids):,} tokens") |
| del texts |
| except Exception as e: |
| print(f" FAIL OSCAR-FR: {e}") |
|
|
| |
| print("4/5 - OpenOrca...") |
| try: |
| ds = load_dataset("Open-Orca/OpenOrca", split="train", streaming=True, trust_remote_code=True) |
| texts = [] |
| count = 0 |
| for x in ds: |
| t = f"System: {x['system_prompt']}\nHuman: {x['question']}\nAssistant: {x['response']}" |
| texts.append(t) |
| count += 1 |
| if count >= 100000: |
| break |
| ids = tokenize_texts(texts, "OpenOrca") |
| all_token_ids.extend(ids) |
| print(f" OK OpenOrca: {len(ids):,} tokens") |
| del texts |
| except Exception as e: |
| print(f" FAIL OpenOrca: {e}") |
|
|
| |
| print("5/5 - Code Python...") |
| try: |
| ds = load_dataset("bigcode/the-stack", data_dir="data/python", split="train", streaming=True, trust_remote_code=True) |
| texts = [] |
| count = 0 |
| for x in ds: |
| texts.append(x["content"]) |
| count += 1 |
| if count >= 50000: |
| break |
| ids = tokenize_texts(texts, "Code-Python") |
| all_token_ids.extend(ids) |
| print(f" OK Code-Python: {len(ids):,} tokens") |
| del texts |
| except Exception as e: |
| print(f" FAIL Code: {e}") |
|
|
| |
| print("Adding local AICL data...") |
| import glob as glob_mod |
| local_ids = [] |
| data_dir = "/root/cognet-1b/data" |
| for filepath in sorted(glob_mod.glob(os.path.join(data_dir, "*.json")) + glob_mod.glob(os.path.join(data_dir, "*.jsonl"))): |
| try: |
| with open(filepath, "r", encoding="utf-8") as f: |
| content = f.read() |
| try: |
| data = json.loads(content) |
| if isinstance(data, dict): |
| for v in data.values(): |
| if isinstance(v, str) and len(v) > 20: |
| local_ids.extend(tokenizer.encode(v)) |
| elif isinstance(v, list): |
| for item in v: |
| if isinstance(item, str) and len(item) > 20: |
| local_ids.extend(tokenizer.encode(item)) |
| elif isinstance(item, dict): |
| for iv in item.values(): |
| if isinstance(iv, str) and len(iv) > 20: |
| local_ids.extend(tokenizer.encode(iv)) |
| except json.JSONDecodeError: |
| if len(content) > 20: |
| local_ids.extend(tokenizer.encode(content)) |
| except: |
| pass |
| all_token_ids.extend(local_ids) |
| print(f" OK Local AICL: {len(local_ids):,} tokens") |
|
|
| |
| print(f"TOTAL TOKENS: {len(all_token_ids):,}") |
| tokens = torch.tensor(all_token_ids, dtype=torch.long) |
| split = int(len(tokens) * 0.95) |
| train_tokens = tokens[:split] |
| val_tokens = tokens[split:] |
| torch.save(train_tokens, os.path.join(DATA_DIR, "train_tokens.pt")) |
| torch.save(val_tokens, os.path.join(DATA_DIR, "val_tokens.pt")) |
| print(f"Train: {len(train_tokens):,} tokens ({len(train_tokens)/1e6:.1f}M)") |
| print(f"Val: {len(val_tokens):,} tokens ({len(val_tokens)/1e6:.1f}M)") |
| print("DATA DOWNLOAD COMPLETE!") |
|
|
