finetuning.py

"""
================================================================================
Tiny-Aya-Global Urdu -> Roman Urdu : Sequential Fine-Tuning Script
v2 (HF, merged) + HBL dataset -> v4 (HF, merged)
================================================================================

Fully offline. No bitsandbytes. No 4-bit quantization. FP16 + LoRA/PEFT.

Reused from the successful Kaggle v2 run (tiny-aya-ft-v2.ipynb):
  - LoRA r=32, alpha=64, target_modules = all 7 linear proj layers
  - max_length=224, "### Instruction / ### Input / ### Response" prompt format
  - Custom sliding-window completion collator (masks everything up to
    "### Response:\n" so loss is only computed on the Roman Urdu tokens)

Changed for this offline A16 16GB run:
  - No BitsAndBytesConfig / 4-bit quant (banned + caused errors before)
  - optim="adamw_torch" instead of "paged_adamw_8bit" (that optimizer requires
    bitsandbytes)
  - HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE forced on, local_files_only=True
    everywhere, no push_to_hub

Run:
    python train_hbl_v4.py
================================================================================
"""

import os

# ------------------------------------------------------------------------
# MUST be set before importing transformers/datasets/huggingface_hub, so
# nothing on this air-gapped box ever attempts an HTTP call.
# ------------------------------------------------------------------------
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_DATASETS_OFFLINE"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import gc
import re
import glob
import random
import unicodedata

import numpy as np
import pandas as pd
import torch
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, PeftModel

random.seed(42)
np.random.seed(42)

# ============================================================================
# CONFIG -- edit these paths for your environment
# ============================================================================

# Local path to the v2 merged HF model (full precision weights + tokenizer),
# copied onto this server. NOT the Ollama GGUF file -- must be the
# safetensors/config/tokenizer HF folder.
MODEL_PATH = "/opt/models/tiny-aya-hbl-v2"

# Local path to the HBL CSV dataset (columns: urdu,roman)
DATASET_PATH = "/opt/data/pairs_clean.csv"

# Dataset has ~137k rows. Adjust this manually to control training set size.
MAX_ROWS = 30000  # <-- adjust manually, dataset has 137k rows

# How many of MAX_ROWS to hold out for eval (taken from the sampled subset)
VAL_ROWS = 1500

# Working directory for checkpoints + LoRA adapter
RUN_DIR = "./tiny-aya-hbl-v4-run"
CHECKPOINT_DIR = os.path.join(RUN_DIR, "checkpoints")
ADAPTER_DIR = os.path.join(RUN_DIR, "final_adapter")

# Final merged full-precision HF model output (model.safetensors, config,
# tokenizer files) -- this is the v4 deliverable.
OUTPUT_DIR = "./tiny-aya-hbl-v4"

# Keep True to match the casing convention used in v2/v3 training
# (previous notebook's clean_roman() uppercased all targets).
# Set False if you want v4 to learn natural-case Roman Urdu instead.
UPPERCASE_ROMAN_TARGETS = True

MAX_SEQ_LEN = 224

# ============================================================================
# 1. DATA LOADING + CLEANING (mirrors notebook's clean_urdu / clean_roman)
# ============================================================================

def clean_urdu(s):
    if not isinstance(s, str) or len(s.strip()) == 0:
        return ""
    s = unicodedata.normalize("NFC", s)
    s = re.sub(r"\s+", " ", s).strip()
    s = re.sub(r",\s*,", ",", s)
    return s


def clean_roman(s):
    if not isinstance(s, str) or len(s.strip()) == 0:
        return ""
    s = unicodedata.normalize("NFC", s)
    s = re.sub(r"\s+", " ", s).strip()
    s = re.sub(r"\s*,\s*", ", ", s)
    s = re.sub(r"\s*\.\s*", ". ", s)
    s = re.sub(r"\s+", " ", s).strip()
    return s.upper() if UPPERCASE_ROMAN_TARGETS else s


def load_hbl_dataset(csv_path, max_rows, val_rows):
    print(f"Loading dataset: {csv_path}")
    df = pd.read_csv(csv_path)

    # Expect columns: urdu, roman
    missing = {"urdu", "roman"} - set(df.columns)
    if missing:
        raise ValueError(f"pairs_clean.csv missing expected columns: {missing}")

    df = df.rename(columns={"urdu": "Urdu_Input", "roman": "Roman_Urdu_Target"})
    df = df.dropna(subset=["Urdu_Input", "Roman_Urdu_Target"])

    print(f"Raw rows: {len(df):,}")

    df["Urdu_Input"] = df["Urdu_Input"].apply(clean_urdu)
    df["Roman_Urdu_Target"] = df["Roman_Urdu_Target"].apply(clean_roman)

    df = df[(df["Urdu_Input"].str.len() > 0) & (df["Roman_Urdu_Target"].str.len() > 0)]
    df = df.drop_duplicates(subset=["Urdu_Input", "Roman_Urdu_Target"])
    df = df[df["Urdu_Input"].str.len() <= 200]
    df = df[df["Roman_Urdu_Target"].str.len() <= 250]
    df = df.reset_index(drop=True)

    print(f"Rows after cleaning/dedup: {len(df):,}")

    if len(df) > max_rows:
        df = df.sample(n=max_rows, random_state=42).reset_index(drop=True)

    print(f"Rows used for this run (MAX_ROWS={max_rows:,}): {len(df):,}")

    val_rows = min(val_rows, max(1, len(df) // 20))
    val_df = df.sample(n=val_rows, random_state=42)
    train_df = df.drop(val_df.index).reset_index(drop=True)
    val_df = val_df.reset_index(drop=True)

    print(f"Train: {len(train_df):,} | Val: {len(val_df):,}")
    return train_df, val_df


def format_example(urdu, roman, eos_token):
    return f"""### Instruction:
Transliterate the following Urdu text into Roman Urdu.
Output ONLY the Roman Urdu. No translation. No explanation.

### Input:
{urdu}

### Response:
{roman}{eos_token}"""


# ============================================================================
# 2. MODEL + TOKENIZER (fp16, no quantization, no bitsandbytes)
# ============================================================================

def load_model_and_tokenizer(model_path):
    print(f"Loading base model (v2) from: {model_path}")

    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        device_map={"": 0},
        trust_remote_code=True,
        local_files_only=True,
    )
    model.config.use_cache = False

    # No prepare_model_for_kbit_training here -- that helper is for
    # quantized (4-bit/8-bit) models only. For plain fp16 we just need
    # gradient checkpointing + input grads enabled manually.
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()

    lora_config = LoraConfig(
        r=32,
        lora_alpha=64,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    return model, tokenizer


# ============================================================================
# 3. TOKENIZATION + COMPLETION-ONLY MASKING COLLATOR
# ============================================================================

def build_tokenized_datasets(train_df, val_df, tokenizer):
    train_texts = [
        format_example(u, r, tokenizer.eos_token)
        for u, r in zip(train_df["Urdu_Input"], train_df["Roman_Urdu_Target"])
    ]
    val_texts = [
        format_example(u, r, tokenizer.eos_token)
        for u, r in zip(val_df["Urdu_Input"], val_df["Roman_Urdu_Target"])
    ]

    train_dataset = Dataset.from_dict({"text": train_texts})
    val_dataset = Dataset.from_dict({"text": val_texts})

    def tokenize_function(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=MAX_SEQ_LEN,
            padding=False,
            return_tensors=None,
        )

    tokenized_train = train_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
    tokenized_val = val_dataset.map(tokenize_function, batched=True, remove_columns=["text"])

    print(f"Tokenized train samples: {len(tokenized_train):,}")
    print(f"Tokenized val samples: {len(tokenized_val):,}")
    print("\nSample formatted example:\n" + train_texts[0])

    return tokenized_train, tokenized_val


def make_completion_collator(tokenizer):
    response_template = "### Response:\n"
    response_ids = tokenizer.encode(response_template, add_special_tokens=False)
    window_size = len(response_ids)

    def custom_completion_collator(features):
        batch = tokenizer.pad(features, return_tensors="pt")
        labels = batch["input_ids"].clone()

        # Mask padding tokens
        labels[batch["attention_mask"] == 0] = -100

        # Mask the prompt: only train on tokens after "### Response:\n"
        for i in range(labels.shape[0]):
            label_seq = labels[i].tolist()
            match_idx = -1
            for j in range(len(label_seq) - window_size + 1):
                if label_seq[j:j + window_size] == response_ids:
                    match_idx = j + window_size
                    break
            if match_idx != -1:
                labels[i, :match_idx] = -100
            else:
                # Safety net: if template wasn't found (e.g. truncation cut
                # it off), don't train on a fully-unmasked prompt-only
                # sequence -- mask the whole thing instead.
                labels[i, :] = -100

        batch["labels"] = labels
        return batch

    return custom_completion_collator


# ============================================================================
# 4. TRAIN
# ============================================================================

def train(model, tokenizer, tokenized_train, tokenized_val):
    os.makedirs(CHECKPOINT_DIR, exist_ok=True)

    training_args = TrainingArguments(
        output_dir=CHECKPOINT_DIR,
        per_device_train_batch_size=4,
        per_device_eval_batch_size=4,
        gradient_accumulation_steps=4,
        num_train_epochs=1,
        learning_rate=2e-4,
        warmup_steps=200,
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=500,
        eval_accumulation_steps=1,
        save_strategy="steps",
        save_steps=500,
        save_total_limit=2,
        load_best_model_at_end=False,
        fp16=True,
        report_to="none",
        dataloader_num_workers=2,
        remove_unused_columns=False,  # required for the custom collator
        optim="adamw_torch",          # NOT paged_adamw_8bit -- that needs bitsandbytes
        push_to_hub=False,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_train,
        eval_dataset=tokenized_val,
        data_collator=make_completion_collator(tokenizer),
    )

    existing_checkpoints = glob.glob(os.path.join(CHECKPOINT_DIR, "checkpoint-*"))
    if existing_checkpoints:
        print(f"Found {len(existing_checkpoints)} existing checkpoints. Resuming...")
        trainer.train(resume_from_checkpoint=True)
    else:
        print("No checkpoints found. Starting fresh training run.")
        trainer.train()

    os.makedirs(ADAPTER_DIR, exist_ok=True)
    model.save_pretrained(ADAPTER_DIR)
    tokenizer.save_pretrained(ADAPTER_DIR)
    print(f"LoRA adapter saved to: {ADAPTER_DIR}")


# ============================================================================
# 5. MERGE ADAPTER INTO BASE (v2) MODEL -> v4 MERGED MODEL
# ============================================================================

def merge_and_save(base_model_path, adapter_path, output_dir):
    # Free GPU memory from the training run before reloading for merge.
    gc.collect()
    torch.cuda.empty_cache()

    print(f"Loading base model (v2) on CPU for safe merging: {base_model_path}")
    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_path,
        torch_dtype=torch.float16,
        device_map="cpu",
        trust_remote_code=True,
        local_files_only=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(base_model_path, local_files_only=True)

    print(f"Merging LoRA adapter from: {adapter_path}")
    merged_model = PeftModel.from_pretrained(base_model, adapter_path)
    merged_model = merged_model.merge_and_unload()

    os.makedirs(output_dir, exist_ok=True)
    merged_model.save_pretrained(output_dir, safe_serialization=True)
    tokenizer.save_pretrained(output_dir)

    print(f"v4 merged model saved to: {output_dir}")
    print("Contents:")
    for f in sorted(os.listdir(output_dir)):
        print(f"  {f}")


# ============================================================================
# 6. QUICK SANITY CHECK (optional, run after merge)
# ============================================================================

def quick_test(merged_model_path, sample_urdu_lines):
    print("\nRunning quick sanity check on merged v4 model...")
    tokenizer = AutoTokenizer.from_pretrained(merged_model_path, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(
        merged_model_path,
        torch_dtype=torch.float16,
        device_map="auto",
        local_files_only=True,
    )
    model.eval()

    for urdu_text in sample_urdu_lines:
        prompt = f"""### Instruction:
Transliterate the following Urdu text into Roman Urdu.
Output ONLY the Roman Urdu. No translation. No explanation.

### Input:
{urdu_text}

### Response:
"""
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=128,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
        result = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
        print(f"\nUrdu:  {urdu_text}")
        print(f"Roman: {result}")


# ============================================================================
# MAIN
# ============================================================================

if __name__ == "__main__":
    print("=" * 70)
    print("Tiny-Aya-Global v2 -> v4 sequential fine-tuning (offline, FP16 LoRA)")
    print("=" * 70)
    print(f"GPU available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

    train_df, val_df = load_hbl_dataset(DATASET_PATH, MAX_ROWS, VAL_ROWS)

    model, tokenizer = load_model_and_tokenizer(MODEL_PATH)

    tokenized_train, tokenized_val = build_tokenized_datasets(train_df, val_df, tokenizer)

    train(model, tokenizer, tokenized_train, tokenized_val)

    # Free the LoRA-wrapped training model before reloading base for merge
    del model
    gc.collect()
    torch.cuda.empty_cache()

    merge_and_save(MODEL_PATH, ADAPTER_DIR, OUTPUT_DIR)

    quick_test(
        OUTPUT_DIR,
        sample_urdu_lines=[
            "آپ کے حبیب بینک لمیٹڈ اکاؤنٹ میں تنخواہ جمع ہو گئی ہے۔",
            "براہ کرم اس دستاویز پر تاریخ درج کریں۔",
        ],
    )

    print("\nDone. v4 model ready at:", OUTPUT_DIR)