| import torch |
| import os |
| import glob |
| import json |
| from tokenizer import Byte_Tokenizer |
|
|
| def load_all_text_data(data_dir): |
| text_data = "" |
| |
| for filepath in glob.glob(os.path.join(data_dir, "*.txt")): |
| with open(filepath, 'r', encoding='utf-8', errors='ignore') as f: |
| text_data += f.read() + "\n" |
| |
| |
| for filepath in glob.glob(os.path.join(data_dir, "*.jsonl")): |
| with open(filepath, 'r', encoding='utf-8', errors='ignore') as f: |
| for line in f: |
| if not line.strip(): |
| continue |
| try: |
| obj = json.loads(line.strip()) |
| if "turns" in obj: |
| for turn in obj["turns"]: |
| text_data += f"{turn.get('speaker', '')}: {turn.get('text', '')}\n" |
| elif "text" in obj: |
| text_data += obj["text"] + "\n" |
| except: |
| pass |
| return text_data |
|
|
| def get_real_data_loader(config, data_dir="data"): |
| print(f"Loading datasets from {data_dir}...") |
| text_data = load_all_text_data(data_dir) |
| if not text_data: |
| raise ValueError(f"No text data found in {data_dir}!") |
| |
| tokenizer = Byte_Tokenizer() |
| print("Tokenizing the entire dataset using Character-level bytes...") |
| tokens = tokenizer.encode(text_data) |
| |
| tensor_data = torch.tensor(tokens, dtype=torch.long) |
| total_len = tensor_data.size(0) |
| |
| chunk_size = config.batch_size * config.seq_len |
| num_chunks = total_len // chunk_size |
| |
| if num_chunks == 0: |
| raise ValueError("Dataset is too small for the configured batch_size and seq_len.") |
| |
| print(f"Dataset loaded: {total_len} tokens, {num_chunks} batches available.") |
| |
| |
| while True: |
| |
| for _ in range(num_chunks): |
| idx = torch.randint(0, total_len - chunk_size - 1, (1,)).item() |
| chunk = tensor_data[idx:idx+chunk_size+1] |
| x = chunk[:-1].view(config.batch_size, config.seq_len) |
| y = chunk[1:].view(config.batch_size, config.seq_len) |
| yield x, y |
|
|
| def save_checkpoint(model, optimizer, step, path="checkpoint.pt"): |
| torch.save({ |
| 'step': step, |
| 'model_state_dict': model.state_dict(), |
| 'optimizer_state_dict': optimizer.state_dict(), |
| }, path) |
| print(f"Checkpoint saved to {path}") |
|
|
| def load_checkpoint(model, optimizer, path="checkpoint.pt"): |
| if os.path.exists(path): |
| checkpoint = torch.load(path) |
| model.load_state_dict(checkpoint['model_state_dict']) |
| optimizer.load_state_dict(checkpoint['optimizer_state_dict']) |
| print(f"Loaded checkpoint from {path} (Step {checkpoint['step']})") |
| return checkpoint['step'] |
| return 0 |
|
|
