Spaces:
Running on Zero
Running on Zero
| 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}") |