# ---
# jupyter:
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
#   kaggle:
#     accelerator: gpu
#     dataSources:
#       - type: competition
#         name: deep-past-initiative-machine-translation
#       - type: dataset
#         name: manwithacat/oracc-akkadian-english-parallel-corpus
#     docker_image: gcr.io/kaggle-gpu-images/python
#     isGpuEnabled: true
#     isInternetEnabled: true
#     language: python
#     sourceType: script
# ---

# %% [markdown]
# # byt5-akkadian-combined
#
# ByT5-base fine-tuned on competition + ORACC data
#
# **Generated from**: `training.toml`
# **Version**: 1.0.0
# **Generated at**: 2025-12-23T20:58:17.097Z

# %%
import os
import sys
import time
import json
import glob
import subprocess
import shutil

os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["TQDM_DISABLE"] = "1"  # Disable progress bars for clean Kaggle logs
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"  # Disable HuggingFace download bars

import torch
import pandas as pd
import numpy as np
from datetime import datetime

if torch.cuda.is_available():
    torch.cuda.empty_cache()

print(f"Python: {sys.version}")
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

# %% [markdown]
# ## Install Dependencies

# %%
packages = ["sentencepiece", "sacrebleu"]  # sacrebleu for chrF metric (no evaluate needed)
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--no-deps"] + packages, check=True)
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "transformers>=4.35.0", "accelerate"], check=True)
print("Packages installed")

# %%
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq,
    EarlyStoppingCallback,
)
import sacrebleu  # Pre-installed on Kaggle, no internet required

# %% [markdown]
# ## Configuration
#
# Auto-generated from `training.toml` - DO NOT EDIT MANUALLY
# To modify, edit the TOML file and regenerate with: akk notebook build training.toml

# %%
# === AUTO-GENERATED CONFIG ===
# Source: training.toml
# Generated: 2025-12-23T20:58:17.097Z
# Version: 1.0.0

CONFIG = {
    # Model
    "model_name": "google/byt5-base",

    # Language codes
    "src_lang": "akk",
    "tgt_lang": "eng",

    # Precision
    "training_precision": "fp32",
    "checkpoint_precision": "fp32",
    "fp16": False,

    # Training
    "num_epochs": 8,
    "batch_size": 2,
    "gradient_accumulation_steps": 4,
    "learning_rate": 0.0002,
    "weight_decay": 0.01,
    "warmup_ratio": 0.06,
    "lr_scheduler_type": "cosine",

    # Sequence lengths
    "max_src_len": 512,
    "max_tgt_len": 256,

    # Evaluation
    "eval_steps": 200,
    "save_steps": 200,
    "logging_steps": 50,
    "predict_with_generate": True,
    "metric_for_best_model": "chrf",
    "greater_is_better": True,

    # Checkpoints
    "save_total_limit": 1,
    "save_only_model": True,  # True = skip optimizer.pt (saves ~4GB per checkpoint)
    "load_best_at_end": True,

    # Early stopping
    "early_stopping_enabled": True,
    "early_stopping_patience": 3,

    # Validation split
    "val_split": 0.1,
    "seed": 42,

    # Generation
    "num_beams": 4,
    "max_new_tokens": 256,
    "repetition_penalty": 1,
    "no_repeat_ngram_size": 0,

    # Platform
    "clear_hf_cache": True,

    # Output
    "output_dir": "./byt5-akkadian",
}

print("Configuration:")
for k, v in CONFIG.items():
    print(f"  {k}: {v}")

# %% [markdown]
# ## Load Dataset

# %%
# Dataset sources (in priority order)
# Note: Use our curated dataset first, competition data doesn't have source/target columns
DATASET_SOURCES = [
    "/kaggle/input/oracc-akkadian-english-parallel-corpus/*.csv",
    "/kaggle/input/deep-past-initiative-machine-translation/train.csv",
]

train_df = None
dataset_type = None

for source in DATASET_SOURCES:
    files = glob.glob(source, recursive=True)
    if files:
        print(f"Using dataset: {files[0]}")
        train_df = pd.read_csv(files[0])
        dataset_type = source.split("/")[-2] if "/" in source else "local"
        break

if train_df is None:
    raise RuntimeError(f"No training data found! Searched: {DATASET_SOURCES}")

print(f"\nTraining samples: {len(train_df)}")
print(f"Columns: {train_df.columns.tolist()}")

# Normalize column names (handle different dataset formats)
column_mapping = {
    'akkadian': 'source',
    'english': 'target',
    'akk': 'source',
    'eng': 'target',
}
train_df = train_df.rename(columns={k: v for k, v in column_mapping.items() if k in train_df.columns})
print(f"Normalized columns: {train_df.columns.tolist()}")

# %%
# Data quality check
print(f"\nData quality:")
print(f"  Null sources: {train_df['source'].isna().sum()}")
print(f"  Null targets: {train_df['target'].isna().sum()}")
print(f"  Avg source length: {train_df['source'].str.len().mean():.1f} chars")
print(f"  Avg target length: {train_df['target'].str.len().mean():.1f} chars")

train_df = train_df.dropna(subset=['source', 'target'])
print(f"  After cleaning: {len(train_df)} samples")

# %% [markdown]
# ## Initialize Model

# %%
print(f"\nLoading model: {CONFIG['model_name']}")
start = time.time()

tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_name"])
model = AutoModelForSeq2SeqLM.from_pretrained(CONFIG["model_name"])

print(f"Model loaded in {time.time()-start:.1f}s")
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

# %%
# Clear HuggingFace cache to save disk space
if CONFIG["clear_hf_cache"]:
    hf_cache = os.path.expanduser("~/.cache/huggingface")
    if os.path.exists(hf_cache):
        cache_size = sum(os.path.getsize(os.path.join(dp, f)) for dp, dn, fn in os.walk(hf_cache) for f in fn)
        print(f"Clearing HF cache: {cache_size / 1e9:.2f} GB")
        shutil.rmtree(hf_cache, ignore_errors=True)
        print("HF cache cleared to save disk space")

# %%
# T5/ByT5 uses task prefix for translation
TASK_PREFIX = "translate Akkadian to English: "
print(f"Using T5 task prefix: {TASK_PREFIX}")

# %% [markdown]
# ## Preprocessing

# %%
def preprocess_function(examples):
    # T5/ByT5: prepend task prefix to inputs
    inputs = [TASK_PREFIX + str(x) for x in examples["source"]]
    targets = [str(x) for x in examples["target"]]

    model_inputs = tokenizer(
        inputs,
        max_length=CONFIG["max_src_len"],
        truncation=True,
        padding=False,
    )

    # Tokenize targets
    labels = tokenizer(
        targets,
        max_length=CONFIG["max_tgt_len"],
        truncation=True,
        padding=False,
    )

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# %%
print("\nPreparing datasets...")

dataset = Dataset.from_pandas(train_df[["source", "target"]])
split = dataset.train_test_split(test_size=CONFIG["val_split"], seed=CONFIG["seed"])
train_dataset = split["train"]
val_dataset = split["test"]

print(f"Train: {len(train_dataset)}, Val: {len(val_dataset)}")

# %%
train_tokenized = train_dataset.map(preprocess_function, batched=True, remove_columns=train_dataset.column_names)
val_tokenized = val_dataset.map(preprocess_function, batched=True, remove_columns=val_dataset.column_names)

print(f"Tokenized - Train: {len(train_tokenized)}, Val: {len(val_tokenized)}")

# %%
# Compute metrics callback for evaluation
# Uses sacrebleu directly (pre-installed on Kaggle) - no internet required
import numpy as np

def compute_metrics(eval_preds):
    preds, labels = eval_preds

    # Replace -100 (ignore index) with pad token before decoding
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)

    # Decode predictions and labels
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    # Strip whitespace
    decoded_preds = [pred.strip() for pred in decoded_preds]
    decoded_labels = [label.strip() for label in decoded_labels]

    # Compute chrF using sacrebleu directly (no internet download needed)
    chrf = sacrebleu.corpus_chrf(decoded_preds, [decoded_labels])

    return {"chrf": chrf.score}

# %% [markdown]
# ## Training Setup

# %%
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, padding=True)

training_args = Seq2SeqTrainingArguments(
    output_dir=CONFIG["output_dir"],
    num_train_epochs=CONFIG["num_epochs"],
    per_device_train_batch_size=CONFIG["batch_size"],
    per_device_eval_batch_size=CONFIG["batch_size"],
    gradient_accumulation_steps=CONFIG["gradient_accumulation_steps"],
    learning_rate=CONFIG["learning_rate"],
    weight_decay=CONFIG["weight_decay"],
    warmup_ratio=CONFIG["warmup_ratio"],
    lr_scheduler_type=CONFIG.get("lr_scheduler_type", "linear"),
    fp16=CONFIG["fp16"],
    eval_strategy="steps",  # Modern API (evaluation_strategy is deprecated)
    eval_steps=CONFIG["eval_steps"],
    save_strategy="steps",
    save_steps=CONFIG["save_steps"],
    save_total_limit=CONFIG["save_total_limit"],
    save_only_model=CONFIG["save_only_model"],  # Skip optimizer.pt to save disk space
    logging_steps=CONFIG["logging_steps"],
    load_best_model_at_end=CONFIG["load_best_at_end"],
    metric_for_best_model=CONFIG["metric_for_best_model"],
    greater_is_better=CONFIG["greater_is_better"],
    predict_with_generate=CONFIG["predict_with_generate"],
    report_to="none",
    seed=CONFIG["seed"],
)

callbacks = []
if CONFIG["early_stopping_enabled"]:
    callbacks.append(EarlyStoppingCallback(early_stopping_patience=CONFIG["early_stopping_patience"]))

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_tokenized,
    eval_dataset=val_tokenized,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    callbacks=callbacks,
)

# %% [markdown]
# ## Train

# %%
print("\n" + "="*60)
print("STARTING TRAINING")
print(f"Config: byt5-akkadian-combined v1.0.0")
print(f"Dataset: {dataset_type} ({len(train_tokenized)} samples)")
print("="*60 + "\n")

train_start = time.time()
train_result = trainer.train()
train_time = time.time() - train_start

print(f"\nTraining completed in {train_time/60:.1f} minutes")
print(f"Final metrics: {train_result.metrics}")

# %% [markdown]
# ## Save Model

# %%
output_dir = CONFIG["output_dir"]
trainer.save_model(output_dir)
tokenizer.save_pretrained(output_dir)

# Save training config
config_path = os.path.join(output_dir, "training_config.json")
with open(config_path, "w") as f:
    json.dump({
        **{k: v for k, v in CONFIG.items() if not callable(v)},
        "dataset_type": dataset_type,
        "train_samples": len(train_tokenized),
        "val_samples": len(val_tokenized),
        "train_time_minutes": train_time / 60,
        "final_loss": train_result.metrics.get("train_loss"),
        "_meta": {
            "name": "byt5-akkadian-combined",
            "version": "1.0.0",
            "source_config": "training.toml",
            "generated_at": "2025-12-23T20:58:17.097Z",
        },
    }, f, indent=2)

# Save length config for inference
length_config = {
    "length_ratio": 1.5,
    "length_slack": 20,
    "global_max": 512,
    "global_min": 32,
}
with open(os.path.join(output_dir, "length_config.json"), "w") as f:
    json.dump(length_config, f)

print(f"\nModel saved to: {output_dir}")
print(f"Files: {os.listdir(output_dir)}")

# %% [markdown]
# ## Upload to Kaggle Model Registry

# %%
# Install kagglehub for model registry upload
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "kagglehub>=0.2.5"], check=True)
import kagglehub

# Upload model to Kaggle registry
kaggle_handle = "manwithacat/akkadian-byt5/transformers/v1"
print(f"\nUploading model to Kaggle: {kaggle_handle}")

try:
    kagglehub.model_upload(
        handle=kaggle_handle,
        local_model_dir=output_dir,
        version_notes=f"byt5-akkadian-combined v1.0.0 - trained on {dataset_type}",
        license_name="Apache 2.0",
    )
    print(f"Model uploaded successfully to: kaggle.com/models/{kaggle_handle.rsplit('/', 2)[0]}")
except Exception as e:
    print(f"Warning: Model upload failed: {e}")
    print("Model is still saved locally and can be uploaded manually")

# %% [markdown]
# ## Quick Evaluation

# %%
print("\n" + "="*60)
print("SAMPLE TRANSLATIONS")
print("="*60)

model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

test_samples = val_dataset.select(range(min(5, len(val_dataset))))

for i, sample in enumerate(test_samples):
    src = sample["source"]
    ref = sample["target"]

    # T5: prepend task prefix
    inputs = tokenizer(TASK_PREFIX + src, return_tensors="pt", max_length=CONFIG["max_src_len"], truncation=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=CONFIG["max_new_tokens"],
            num_beams=CONFIG["num_beams"],
            repetition_penalty=CONFIG["repetition_penalty"],
            no_repeat_ngram_size=CONFIG["no_repeat_ngram_size"],
        )

    pred = tokenizer.decode(outputs[0], skip_special_tokens=True)

    print(f"\n[{i}] SRC: {src[:80]}...")
    print(f"    REF: {ref[:80]}...")
    print(f"    OUT: {pred[:80]}...")

# %% [markdown]
# ## Generate Competition Submission

# %%
# Load test data
test_files = glob.glob("/kaggle/input/deep-past-initiative-machine-translation/test.csv", recursive=True)
print(f"Found test files: {test_files}")

if not test_files:
    print("WARNING: test.csv not found - skipping submission generation")
    test_df = None
else:
    test_df = pd.read_csv(test_files[0])
    print(f"Test samples: {len(test_df)}")
    print(test_df.head())

# %%
# === OPTIMIZED BATCH INFERENCE ===
# Key optimizations for competition scoring:
# 1. FP16 inference (2x faster, less VRAM)
# 2. Larger batch size (no gradients = more headroom)
# 3. Greedy decoding (num_beams=1, much faster than beam search)
# 4. Sorted batching (reduces padding overhead)
# 5. torch.inference_mode (faster than no_grad)
# 6. torch.compile (JIT compilation, PyTorch 2.0+)

INFERENCE_BATCH_SIZE = 16  # Larger batches for inference (no gradients)
INFERENCE_NUM_BEAMS = 1    # Greedy decoding (fastest) - set to 2 for slight quality boost

# Convert to FP16 for faster inference
print("\nOptimizing model for inference...")
model = model.half()  # FP16 inference
model.eval()

# Try torch.compile for additional speedup (PyTorch 2.0+)
try:
    if hasattr(torch, 'compile'):
        model = torch.compile(model, mode="reduce-overhead")
        print("  torch.compile enabled")
except Exception as e:
    print(f"  torch.compile not available: {e}")

def translate_batch_optimized(texts, batch_size=INFERENCE_BATCH_SIZE):
    """
    Optimized batch translation with:
    - Sorted batching by length (reduces padding)
    - FP16 inference
    - Greedy or low-beam decoding
    """
    if not texts:
        return []

    # Sort by length for efficient batching (less padding waste)
    indexed_texts = [(i, t) for i, t in enumerate(texts)]
    indexed_texts.sort(key=lambda x: len(x[1]))

    all_translations = [None] * len(texts)

    with torch.inference_mode():  # Faster than no_grad
        for batch_start in range(0, len(indexed_texts), batch_size):
            batch = indexed_texts[batch_start:batch_start + batch_size]
            batch_indices = [x[0] for x in batch]
            batch_texts = [TASK_PREFIX + x[1] for x in batch]

            # Tokenize with dynamic padding (pad to longest in batch)
            inputs = tokenizer(
                batch_texts,
                max_length=CONFIG["max_src_len"],
                truncation=True,
                padding="longest",  # Dynamic padding per batch
                return_tensors="pt",
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}

            outputs = model.generate(
                **inputs,
                max_new_tokens=CONFIG["max_new_tokens"],
                num_beams=INFERENCE_NUM_BEAMS,
                do_sample=False,  # Deterministic
                use_cache=True,   # KV cache for autoregressive generation
                early_stopping=True if INFERENCE_NUM_BEAMS > 1 else False,
            )

            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)

            # Map back to original order
            for idx, translation in zip(batch_indices, decoded):
                all_translations[idx] = translation

    return all_translations

# %%
if test_df is not None:
    # Detect source column
    source_col = "transliteration" if "transliteration" in test_df.columns else "source" if "source" in test_df.columns else test_df.columns[1]
    sources = test_df[source_col].tolist()

    print(f"\nTranslating {len(sources)} test samples using column: {source_col}")
    print(f"  Batch size: {INFERENCE_BATCH_SIZE}, Beams: {INFERENCE_NUM_BEAMS}, Precision: FP16")

    import time as _time
    _start = _time.time()

    # Single optimized call for all translations
    translations = translate_batch_optimized(sources)

    _elapsed = _time.time() - _start
    print(f"Generated {len(translations)} translations in {_elapsed:.1f}s ({len(sources)/_elapsed:.1f} samples/sec)")

    # Create submission
    id_col = "id" if "id" in test_df.columns else "id"
    submission = pd.DataFrame({
        id_col: test_df[id_col] if id_col in test_df.columns else range(len(translations)),
        "translation": translations
    })

    print("\nSubmission preview:")
    print(submission.head(10))

    submission.to_csv("submission.csv", index=False)
    print("\nSaved submission.csv")

    # Show sample translations
    print("\n=== Sample Test Translations ===")
    for i in range(min(5, len(sources))):
        print(f"\n[{i}]")
        print(f"SRC: {sources[i][:100]}...")
        print(f"OUT: {translations[i][:100]}...")

# %%
print("\n" + "="*60)
print("TRAINING COMPLETE")
print(f"Config: byt5-akkadian-combined v1.0.0")
print(f"Dataset: {dataset_type}")
print(f"Samples: {len(train_tokenized)} train, {len(val_tokenized)} val")
print(f"Time: {train_time/60:.1f} minutes")
print(f"Model saved to: {output_dir}")
print("="*60)