""" Kaggle GPU Training Script: ByT5-base for Akkadian-to-English Translation Optimized for GPU (T4/P100) on Kaggle with checkpoint resilience. Features (v3 - GPU): - OptimizedPreprocessor: gap normalization + character corrections - SmartChunker: split long pairs at natural boundaries - VisualizationCallback: print sample translations during evaluation - Checkpoint resilience: auto-resume, frequent saves, early stopping - Sentence alignment + bidirectional augmentation Usage on Kaggle: 1. Enable GPU Accelerator (T4 x2 or P100) in Notebook Settings. 2. Run this script. """ import gc import os import re import logging import sys import subprocess from typing import List import time # --------------------------------------------------------------------------- # Install missing libraries (Kaggle specific) # --------------------------------------------------------------------------- def install_dependencies(): packages = ["evaluate", "sacrebleu", "sentencepiece", "scikit-learn"] print("Installing libraries for GPU environment...") subprocess.check_call( [sys.executable, "-m", "pip", "install", "--quiet"] + packages ) # Call immediately install_dependencies() import numpy as np import pandas as pd import torch import evaluate from datasets import Dataset from transformers import ( AutoModelForSeq2SeqLM, AutoTokenizer, DataCollatorForSeq2Seq, EarlyStoppingCallback, Seq2SeqTrainer, Seq2SeqTrainingArguments, TrainerCallback, ) # --------------------------------------------------------------------------- # Config (GPU / ByT5-base) # --------------------------------------------------------------------------- class Config: SEED = 42 MODEL_NAME = "google/byt5-base" # UPGRADED from small # GPU T4 16GB: byt5-base fits with batch_size=1 + gradient_checkpointing PER_DEVICE_BATCH_SIZE = 1 GRADIENT_ACCUMULATION = 8 # Effective batch size = 8 MAX_LENGTH = 256 # Akkadian sentences are usually short. 256 is safe and fast. EPOCHS = 15 # With early stopping, we won't necessarily run all LEARNING_RATE = 5e-5 LABEL_SMOOTHING = 0.1 # Checkpoint resilience SAVE_STEPS = 200 # Save every 200 steps (~frequent enough to not lose progress) EARLY_STOPPING_PATIENCE = 3 # Stop if no improvement for 3 evals # Paths (Kaggle or Local) if os.path.exists("/kaggle"): DATA_DIR = "/kaggle/input/deep-past-initiative-machine-translation" OUTPUT_DIR = "/kaggle/working/models/byt5-base-akkadian" LOG_DIR = "/kaggle/working/logs" else: DATA_DIR = "data/raw" OUTPUT_DIR = "results/kaggle_gpu_dryrun" LOG_DIR = "logs/kaggle_gpu_dryrun" # Validation VAL_SIZE = 0.05 NUM_BEAMS_EVAL = 4 # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[logging.StreamHandler(sys.stdout)] ) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # OptimizedPreprocessor (ported from inference_mbr.py) # --------------------------------------------------------------------------- class OptimizedPreprocessor: """Normalize Akkadian transliteration text before tokenization.""" def __init__(self): self.patterns = { "big_gap": re.compile(r"(\.{3,}|…+|……)"), "small_gap": re.compile(r"(xx+|\s+x\s+)"), } def preprocess(self, text: str) -> str: if pd.isna(text): return "" cleaned = str(text) # Gap normalization cleaned = self.patterns["big_gap"].sub("", cleaned) cleaned = self.patterns["small_gap"].sub("", cleaned) # Critical character corrections (match inference pipeline) cleaned = cleaned.replace("aa", "ša") cleaned = cleaned.replace("kà-ni-ia", "kà-ni-iš") cleaned = cleaned.replace("mup", "ṭup") cleaned = cleaned.replace("aí", "ší") cleaned = cleaned.replace("u\u201e", "u₄") cleaned = cleaned.replace("ia-tù", "iš-tù") cleaned = cleaned.replace("aé", "šé") return cleaned def preprocess_series(self, series: pd.Series) -> pd.Series: """Vectorized preprocessing for a pandas Series.""" s = series.fillna("").astype(str) s = s.str.replace(self.patterns["big_gap"], "", regex=True) s = s.str.replace(self.patterns["small_gap"], "", regex=True) s = s.str.replace("aa", "ša", regex=False) s = s.str.replace("kà-ni-ia", "kà-ni-iš", regex=False) s = s.str.replace("mup", "ṭup", regex=False) s = s.str.replace("aí", "ší", regex=False) s = s.str.replace("u\u201e", "u₄", regex=False) s = s.str.replace("ia-tù", "iš-tù", regex=False) s = s.str.replace("aé", "šé", regex=False) return s # --------------------------------------------------------------------------- # SmartChunker (ported from inference_mbr.py) # --------------------------------------------------------------------------- class SmartChunker: """Split long text at natural boundaries (period, semicolon, comma, space).""" def __init__(self, max_len: int = 400, overlap: int = 50): self.max_len = max_len self.overlap = overlap self.delimiters = [ re.compile(r"(\.\s+)"), re.compile(r"(;\s+)"), re.compile(r"(,\s+)"), re.compile(r"(\s+)"), ] def chunk_text(self, text: str) -> List[str]: if pd.isna(text): return [""] text = str(text) if len(text) <= self.max_len: return [text] chunks = [] remaining = text while len(remaining) > self.max_len: split_idx = -1 search_start = self.max_len - self.overlap search_end = self.max_len zone = remaining[search_start:search_end] found = False for pattern in self.delimiters: match = list(pattern.finditer(zone)) if match: last_match = match[-1] split_idx = search_start + last_match.end() found = True break if not found: split_idx = self.max_len chunks.append(remaining[:split_idx].strip()) remaining = remaining[split_idx:].strip() if remaining: chunks.append(remaining) return chunks def smart_chunk_pairs(df: pd.DataFrame, chunker: SmartChunker) -> pd.DataFrame: """Split long source+target pairs using SmartChunker. For training, both source and target must be split together. Strategy: chunk the source, then proportionally split the target by word count to match the number of source chunks. """ new_rows = [] split_count = 0 for _, row in df.iterrows(): src = str(row["transliteration"]) tgt = str(row["translation"]) src_chunks = chunker.chunk_text(src) if len(src_chunks) == 1: # No splitting needed new_rows.append({"transliteration": src, "translation": tgt}) else: # Split target proportionally to match source chunks split_count += 1 n = len(src_chunks) tgt_words = tgt.split() total_tgt = len(tgt_words) # Distribute target words proportionally based on source chunk lengths src_lens = [len(c) for c in src_chunks] total_src_len = sum(src_lens) tgt_chunks = [] offset = 0 for i, sl in enumerate(src_lens): if i == n - 1: # Last chunk gets all remaining target words tgt_chunk = " ".join(tgt_words[offset:]) else: count = max(1, round(total_tgt * sl / total_src_len)) tgt_chunk = " ".join(tgt_words[offset:offset + count]) offset += count tgt_chunks.append(tgt_chunk) for sc, tc in zip(src_chunks, tgt_chunks): if sc.strip() and tc.strip(): new_rows.append({"transliteration": sc, "translation": tc}) if split_count > 0: logger.info(f" SmartChunker split {split_count} long pairs into {len(new_rows)} total pairs.") return pd.DataFrame(new_rows) # --------------------------------------------------------------------------- # VisualizationCallback # --------------------------------------------------------------------------- class VisualizationCallback(TrainerCallback): """Print sample translations at each evaluation to monitor quality.""" def __init__(self, tokenizer, test_samples, max_length=256): self.tokenizer = tokenizer self.test_samples = test_samples self.max_length = max_length def on_evaluate(self, args, state, control, **kwargs): model = kwargs['model'] model.eval() print(f"\n--- Generation Check at Step {state.global_step} ---") for text in self.test_samples: input_text = text if not text.startswith("translate Akkadian to English:"): input_text = "translate Akkadian to English: " + text inputs = self.tokenizer( input_text, return_tensors="pt", truncation=True, max_length=self.max_length ).to(model.device) with torch.no_grad(): outputs = model.generate( **inputs, max_length=self.max_length, num_beams=4, early_stopping=True, ) decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True) print(f"Source: {text[:100]}{'...' if len(text) > 100 else ''}") print(f"Gen: {decoded}\n") model.train() # --------------------------------------------------------------------------- # Seed # --------------------------------------------------------------------------- def set_seed(seed: int) -> None: np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) # --------------------------------------------------------------------------- # Data Processing (Same logic as baseline) # --------------------------------------------------------------------------- def sentence_align(df: pd.DataFrame) -> pd.DataFrame: """Split document-level pairs into sentence-level pairs.""" aligned = [] for _, row in df.iterrows(): src = str(row["transliteration"]) tgt = str(row["translation"]) # Split English by sentence-ending punctuation tgt_sents = [s.strip() for s in re.split(r"(?<=[.!?])\s+", tgt) if s.strip()] src_words = src.split() n_sents = len(tgt_sents) if n_sents > 1 and len(src_words) >= n_sents * 2: tgt_lens = [len(s.split()) for s in tgt_sents] total_tgt_len = sum(tgt_lens) if total_tgt_len > 0: src_counts = [] remaining = len(src_words) for i, tl in enumerate(tgt_lens): if i == n_sents - 1: count = remaining else: count = max(1, round(len(src_words) * tl / total_tgt_len)) remaining -= count src_counts.append(count) offset = 0 valid = True pairs = [] for count, tgt_sent in zip(src_counts, tgt_sents): src_chunk = " ".join(src_words[offset : offset + count]) offset += count if len(src_chunk) > 3 and len(tgt_sent) > 3: pairs.append({"transliteration": src_chunk, "translation": tgt_sent}) else: valid = False break if valid and len(pairs) == n_sents: aligned.extend(pairs) continue aligned.append({"transliteration": src, "translation": tgt}) return pd.DataFrame(aligned) def make_bidirectional(dataset_split: Dataset) -> Dataset: """Akk→Eng + Eng→Akk""" df = dataset_split.to_pandas() # Forward df_fwd = df.copy() df_fwd["input_text"] = "translate Akkadian to English: " + df_fwd["transliteration"].astype(str) df_fwd["target_text"] = df_fwd["translation"].astype(str) # Backward df_bwd = df.copy() df_bwd["input_text"] = "translate English to Akkadian: " + df_bwd["translation"].astype(str) df_bwd["target_text"] = df_bwd["transliteration"].astype(str) combined = pd.concat([df_fwd, df_bwd], ignore_index=True) combined = combined.sample(frac=1, random_state=Config.SEED).reset_index(drop=True) return Dataset.from_pandas(combined) def make_unidirectional(dataset_split: Dataset) -> Dataset: """Validation: Akk→Eng only""" df = dataset_split.to_pandas() df["input_text"] = "translate Akkadian to English: " + df["transliteration"].astype(str) df["target_text"] = df["translation"].astype(str) return Dataset.from_pandas(df) # --------------------------------------------------------------------------- # Metrics # --------------------------------------------------------------------------- def compute_metrics(eval_preds, tokenizer, metric_chrf, metric_bleu): preds, labels = eval_preds if isinstance(preds, tuple): preds = preds[0] if hasattr(preds, "ndim") and preds.ndim == 3: preds = np.argmax(preds, axis=-1) preds = preds.astype(np.int64) preds = np.where(preds < 0, tokenizer.pad_token_id, preds) preds = np.clip(preds, 0, tokenizer.vocab_size - 1) decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True) labels = np.where(labels != -100, labels, tokenizer.pad_token_id) decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) chrf = metric_chrf.compute(predictions=decoded_preds, references=decoded_labels)["score"] bleu = metric_bleu.compute(predictions=decoded_preds, references=[[x] for x in decoded_labels])["score"] geo_mean = (chrf * bleu) ** 0.5 if chrf > 0 and bleu > 0 else 0.0 return {"chrf": chrf, "bleu": bleu, "geo_mean": geo_mean} # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> None: set_seed(Config.SEED) os.makedirs(Config.OUTPUT_DIR, exist_ok=True) os.makedirs(Config.LOG_DIR, exist_ok=True) logger.info("Starting Kaggle TPU Training Script...") # 1. Load Data logger.info("Scanning for train.csv...") found_train_path = None # Kaggle specific path vs Local fallback possible_paths = [ "/kaggle/input/deep-past-initiative-machine-translation/train.csv", "data/raw/train.csv", # Local fallback "train.csv" ] for p in possible_paths: if os.path.exists(p): found_train_path = p break if not found_train_path: # Deep search in /kaggle/input if likely there if os.path.isdir("/kaggle/input"): for root, dirs, files in os.walk("/kaggle/input"): if "train.csv" in files: found_train_path = os.path.join(root, "train.csv") break if not found_train_path: logger.error("Could not find train.csv!") return logger.info(f"Loading data from: {found_train_path}") train_df = pd.read_csv(found_train_path) # ---- Preprocessing: normalize Akkadian text ---- preprocessor = OptimizedPreprocessor() logger.info("Applying OptimizedPreprocessor to source text...") train_df["transliteration"] = preprocessor.preprocess_series(train_df["transliteration"]) logger.info(f" Preprocessed {len(train_df)} rows.") # Sentence Align train_aligned = sentence_align(train_df) logger.info(f"Aligned pairs: {len(train_aligned)}") # ---- SmartChunker: split long pairs ---- chunker = SmartChunker(max_len=400, overlap=50) train_chunked = smart_chunk_pairs(train_aligned, chunker) logger.info(f"After SmartChunking: {len(train_chunked)} pairs (was {len(train_aligned)})") # Split dataset = Dataset.from_pandas(train_chunked) split = dataset.train_test_split(test_size=Config.VAL_SIZE, seed=Config.SEED) train_dataset = make_bidirectional(split["train"]) val_dataset = make_unidirectional(split["test"]) # 2. Model & Tokenizer logger.info(f"Loading Model: {Config.MODEL_NAME}") tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME) model = AutoModelForSeq2SeqLM.from_pretrained(Config.MODEL_NAME) # Preprocessing def preprocess(examples): inputs = [str(x) for x in examples["input_text"]] targets = [str(x) for x in examples["target_text"]] model_inputs = tokenizer(inputs, max_length=Config.MAX_LENGTH, truncation=True, padding="max_length") labels = tokenizer(targets, max_length=Config.MAX_LENGTH, truncation=True, padding="max_length") model_inputs["labels"] = labels["input_ids"] return model_inputs logger.info("Tokenizing...") tokenized_train = train_dataset.map(preprocess, batched=True) tokenized_val = val_dataset.map(preprocess, batched=True) # 3. Training Args (GPU Optimized + Checkpoint Resilience) is_kaggle = os.path.exists("/kaggle") training_args = Seq2SeqTrainingArguments( output_dir=Config.OUTPUT_DIR, # --- Evaluation & Saving (checkpoint resilience) --- eval_strategy="steps", eval_steps=Config.SAVE_STEPS, save_strategy="steps", save_steps=Config.SAVE_STEPS, save_total_limit=3, # Keep last 3 checkpoints load_best_model_at_end=True, # Auto-load best checkpoint at end metric_for_best_model="eval_loss", # Use eval_loss to pick best greater_is_better=False, # --- Hyperparams --- learning_rate=Config.LEARNING_RATE, optim="adafactor", # Memory efficient optimizer label_smoothing_factor=Config.LABEL_SMOOTHING, per_device_train_batch_size=Config.PER_DEVICE_BATCH_SIZE if is_kaggle else 1, per_device_eval_batch_size=Config.PER_DEVICE_BATCH_SIZE if is_kaggle else 1, gradient_accumulation_steps=Config.GRADIENT_ACCUMULATION, num_train_epochs=Config.EPOCHS if is_kaggle else 1, weight_decay=0.01, # --- Generation (for eval) --- predict_with_generate=True, generation_max_length=Config.MAX_LENGTH, generation_num_beams=Config.NUM_BEAMS_EVAL, # --- GPU Optimization --- fp16=False, # ByT5 produces NaN with fp16 bf16=False, # Safe default; T4 doesn't support bf16 gradient_checkpointing=True, # Critical for byt5-base on T4 16GB dataloader_num_workers=2 if is_kaggle else 0, # --- Logging --- logging_steps=50, logging_dir=Config.LOG_DIR, report_to="none", seed=Config.SEED, ) # Speed up local test by limiting data if not is_kaggle: logger.info("Local mode: Using tiny dataset subsample.") tokenized_train = tokenized_train.select(range(min(10, len(tokenized_train)))) tokenized_val = tokenized_val.select(range(min(10, len(tokenized_val)))) data_collator = DataCollatorForSeq2Seq(tokenizer, model=model) metric_chrf = evaluate.load("chrf") metric_bleu = evaluate.load("sacrebleu") trainer = Seq2SeqTrainer( model=model, args=training_args, train_dataset=tokenized_train, eval_dataset=tokenized_val, data_collator=data_collator, processing_class=tokenizer, compute_metrics=lambda eval_preds: compute_metrics(eval_preds, tokenizer, metric_chrf, metric_bleu), ) # ---- Early Stopping ---- trainer.add_callback(EarlyStoppingCallback(early_stopping_patience=Config.EARLY_STOPPING_PATIENCE)) logger.info(f"Added EarlyStoppingCallback (patience={Config.EARLY_STOPPING_PATIENCE})") # ---- VisualizationCallback ---- try: viz_samples = val_dataset.select(range(min(5, len(val_dataset))))["transliteration"] trainer.add_callback(VisualizationCallback(tokenizer, viz_samples, max_length=Config.MAX_LENGTH)) logger.info(f"Added VisualizationCallback with {len(viz_samples)} samples.") except Exception as e: logger.warning(f"Failed to add VisualizationCallback: {e}") # ---- Auto-resume from checkpoint ---- resume_checkpoint = None if os.path.isdir(Config.OUTPUT_DIR): checkpoints = [d for d in os.listdir(Config.OUTPUT_DIR) if d.startswith("checkpoint-")] if checkpoints: latest = sorted(checkpoints, key=lambda x: int(x.split("-")[-1]))[-1] resume_checkpoint = os.path.join(Config.OUTPUT_DIR, latest) logger.info(f"Resuming from checkpoint: {resume_checkpoint}") logger.info("Starting Training...") trainer.train(resume_from_checkpoint=resume_checkpoint) # ---- Save final model ---- logger.info("Saving Model...") trainer.save_model(Config.OUTPUT_DIR) tokenizer.save_pretrained(Config.OUTPUT_DIR) # ---- Final evaluation ---- logger.info("Running final evaluation...") results = trainer.evaluate() logger.info(f"Final results: BLEU={results.get('eval_bleu', 'N/A')}, " f"chrF={results.get('eval_chrf', 'N/A')}, " f"GeoMean={results.get('eval_geo_mean', 'N/A')}") logger.info("Training Complete.") if __name__ == "__main__": main()