Spaces:
Running on Zero
Running on Zero
File size: 3,891 Bytes
1717c0c | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | from pathlib import Path
import torch
import json
from src.infra.dataset.util import split_tokens, get_batch
from src.infra.tokenizer import CharacterTokenizer
from src.infra.dataset.answering import QuestionAnsweringDataset
from torch.utils.data import random_split
from src.task.config import ModelConfig,TrainingConfig
from src.task.model import TinyAnsweringModel
from src.task.training.question_answering import train
from src.task.training.util import TrainingLogger, count_parameters, separator
logger = TrainingLogger(
"logs/question_answering.log"
)
checkpoint = "tiny_answering.pt"
# Configs
separator("Configuration")
torch.manual_seed(42)
# Configs
model_config = ModelConfig()
training_config = TrainingConfig()
# Overrides
model_config.context_length=512
training_config.num_steps = 300
training_config.eval_interval = 10
logger.log(model_config)
logger.log(training_config)
# Load the corpus
separator("Corpus")
torch.manual_seed(42)
SQUAD_PATH = Path("data/squad/dev-v1.1.json")
QA_TOKENIZER_PATH = Path("checkpoints/qa_tokenizer.json")
print(f"Squad file: {SQUAD_PATH}")
print(f"Tokenizer file: {QA_TOKENIZER_PATH}")
if not SQUAD_PATH.exists():
raise FileNotFoundError(f"Squad file not found: {SQUAD_PATH}")
if not QA_TOKENIZER_PATH.exists():
raise FileNotFoundError(f"Tokenizer file not found: {QA_TOKENIZER_PATH}")
tokenizer = CharacterTokenizer.load(QA_TOKENIZER_PATH)
vocab_size = tokenizer.vocab_size
with SQUAD_PATH.open("r", encoding="utf-8") as file:
squad = json.load(file)
# Tokens
records = []
article = next(
article
for article in squad["data"]
if article["title"] == "Kenya"
)
# for index, article in enumerate(squad["data"]): #when we need all the data
for paragraph in article["paragraphs"] :
context = paragraph["context"]
for qa in paragraph["qas"]:
for answer in qa["answers"]:
records.append(
{
"question": qa["question"],
"context": context,
"answer": answer["text"],
"answer_start": answer["answer_start"],
}
)
print(f"Number of records: {len(records)}")
print("\n=== Sample record ===")
question = records[0]["question"]
context = records[0]["context"]
answer_text = records[0]["answer"]
answer_start = records[0]["answer_start"]
answer_end = answer_start + len(answer_text)
print("Context:", context)
print("Question:", question)
print("Answer:", answer_text)
print("Answer Start:", answer_start)
print("Answer End:", answer_end)
separator("Dataset")
dataset = QuestionAnsweringDataset(
records=records,
tokenizer=tokenizer,
max_length=model_config.context_length, #to manage truncation and padding
)
n = len(dataset)
train_size = int(0.90 * n)
val_size = int(0.05 * n)
test_size = n - train_size - val_size
train_dataset, val_dataset, test_dataset = random_split(
dataset,
[train_size, val_size, test_size],
generator=torch.Generator().manual_seed(42),
)
logger.log(f"Train examples: {len(train_dataset):,}")
logger.log(f"Val examples: {len(val_dataset):,}")
logger.log(f"Test examples: {len(test_dataset):,}")
# Models
separator("Model")
model = TinyAnsweringModel(
config=model_config,
vocab_size=tokenizer.vocab_size,
)
logger.log(model)
# Train........
separator("Training")
model = train(
model=model,
train_dataset=train_dataset,
val_dataset=val_dataset,
training_config=training_config,
)
logger.log(f"Parameters: {count_parameters(model):,}")
separator("Saving")
checkpoint_dir = Path("checkpoints")
checkpoint_dir.mkdir(parents=True, exist_ok=True)
checkpoint_path = (
checkpoint_dir / checkpoint
)
torch.save(
model.state_dict(),
checkpoint_path,
)
print(f"Checkpoint saved to: {checkpoint_path}") |