Spaces:
Running on Zero
Running on Zero
File size: 3,084 Bytes
98d1e65 5bdeead 98d1e65 5bdeead 98d1e65 5bdeead 98d1e65 5bdeead 137f7ef 5bdeead 98d1e65 137f7ef 98d1e65 | 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 | from pathlib import Path
import torch
from src.infra.dataset.classification import TextClassificationDataset
from src.infra.dataset.util import split_tokens, get_batch, build_label, build_text
from src.infra.tokenizer import CharacterTokenizer
from src.infra.dataset.autoregressive import AutoRegressiveDataset
from src.task.config import ModelConfig,TrainingConfig
from src.task.model import TinyTransformerLM, TinyClassifier
from src.task.training.classification import train
import pandas as pd
import numpy as np
from src.task.training.util import TrainingLogger, count_parameters, separator
# Custom Logger
logger = TrainingLogger(
"logs/text_classifier.log"
)
# Configs
separator("Configuration")
torch.manual_seed(42)
# Configs
model_config = ModelConfig()
training_config = TrainingConfig()
# Overrides
model_config.context_length = 384
model_config.number_of_heads = 2
training_config.num_steps = 10
training_config.eval_interval = 1
training_config.num_train_samples = None
logger.log(model_config)
logger.log(training_config)
# Load the corpus
separator("Corpus")
TRAIN_PATH = Path("data/ag-news/train.csv")
TEST_PATH = Path("data/ag-news/test.csv")
checkpoint = "mh_tiny_text_classifier.pt"
if not TRAIN_PATH.exists():
raise FileNotFoundError(f"Training file not found: {TRAIN_PATH}")
if not TEST_PATH.exists():
raise FileNotFoundError(f"Test file not found: {TEST_PATH}")
separator("Load corpus")
train_corpus = pd.read_csv(TRAIN_PATH)
if training_config.num_train_samples is not None:
train_corpus = train_corpus.head(training_config.num_train_samples)
test_corpus = pd.read_csv(TEST_PATH)
train_text = build_text(train_corpus)
test_text = build_text(test_corpus)
train_labels = build_label(train_corpus)
test_labels = build_label(test_corpus)
# Tokens
separator("Tokenization")
tokenizer = CharacterTokenizer("".join(train_text))
CLASS_TOKENIZER_PATH = Path("checkpoints/ag_news_tokenizer.json")
tokenizer.save(CLASS_TOKENIZER_PATH)
vocab_size = tokenizer.vocab_size
print(f"Vocab size: {vocab_size}")
# Datasets
separator("Datasets")
train_dataset = TextClassificationDataset(
texts=train_text,
labels=train_labels,
tokenizer=tokenizer,
max_length=model_config.context_length
)
val_dataset = TextClassificationDataset(
texts=test_text,
labels=test_labels,
tokenizer=tokenizer,
max_length=model_config.context_length
)
print(f"Train examples: {len(train_dataset):,}")
# Models
separator("Transformer")
model = TinyClassifier(
config=model_config,
vocab_size=tokenizer.vocab_size,
num_classes=4,
)
print(model)
# Train........
separator("Training")
train(
model=model,
train_dataset=train_dataset,
val_dataset=val_dataset,
training_config=training_config,
)
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}") |