File size: 2,896 Bytes
dbd41fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import torch
import os
import glob
import json
from tokenizer import Byte_Tokenizer
def load_all_text_data(data_dir):
text_data = ""
# Load .txt files
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"
# Load .jsonl files
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.")
# Infinite generator
while True:
# Create random offset to yield different batches
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
|