Create finetune.py

finetune.py

@@ -1,38 +1,32 @@
+#!/usr/bin/env python3
 """
-============================================================
- Urdu -> Roman Urdu Fine-Tuning Script (CPU / Offline / HBL)
-============================================================
-Converted from tiny-aya-ft-v2.ipynb for offline, CPU-only,
-air-gapped execution on the HBL server.
-
-Key differences from the original notebook:
-  - No internet access anywhere (HF_HUB_OFFLINE=1)
-  - No bitsandbytes / 4-bit quantization (CPU doesn't support it)
-  - No push_to_hub / HF login (nothing leaves this machine)
-  - Loads YOUR local data.csv (columns: urdu,roman) instead of
-    the HF Roman-Urdu-Parl-split dataset + synthetic generator
-  - Loads your already-finetuned model from a local folder
-    instead of downloading CohereLabs/tiny-aya-global
-
-Run with:
-    C:\\Python311\\python.exe finetune_hbl.py
-============================================================
+HBL GPU Server - Stage 2 Fine-tuning Script
+============================================
+Further fine-tunes your existing Urdu -> Roman Urdu transliteration model
+using ONLY pairs_clean.csv (real NADRA log data), then merges the new
+LoRA adapter into your existing fine-tuned model and saves a brand new
+merged model locally.
+
+FULLY OFFLINE - no internet / Hugging Face Hub access is used anywhere.
+Install packages on a machine with wifi, transfer the wheel files here,
+then `pip install --no-index --find-links=<folder> -r requirements.txt`.
+See requirements.txt for the package list and offline-install notes.
 """
 
 import os
 
-# ------------------------------------------------------------
-# MUST be set before importing transformers/datasets/huggingface_hub
-# ------------------------------------------------------------
+# --- Hard block any accidental network calls (safe on an air-gapped box) ---
 os.environ["HF_HUB_OFFLINE"] = "1"
 os.environ["TRANSFORMERS_OFFLINE"] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
 import gc
 import re
-import random
 import unicodedata
 import pandas as pd
 import torch
+from sklearn.model_selection import train_test_split
 from datasets import Dataset
 from transformers import (
     AutoTokenizer,
@@ -40,72 +34,81 @@ from transformers import (
     TrainingArguments,
     Trainer,
 )
-from peft import LoraConfig, get_peft_model, PeftModel
-
-random.seed(42)
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    PeftModel,
+)
 
 # ============================================================
-# CONFIG - edit these paths/values for your setup
+# CONFIG - EDIT THESE FOR YOUR HBL GPU SERVER
 # ============================================================
 
-# Path to your ALREADY fine-tuned tiny-aya-global model folder
-# (the one you copied over earlier, confirmed working)
-MODEL_PATH = r"D:\urdu_model\your_finetuned_tiny_aya_model"
-
-# Path to your HBL dataset: CSV with header "urdu,roman"
-DATA_PATH = r"D:\urdu_model\data.csv"
-
-# Where to save the new LoRA adapter + merged model
-OUTPUT_DIR = r"D:\urdu_model\hbl_finetune_output"
-
-# --------------------------------------------------------------------
-# ROW LIMIT - IMPORTANT for CPU training
-# --------------------------------------------------------------------
-# CPU fine-tuning is roughly 20-50x slower than a T4 GPU. As a rough
-# guide on a modern multi-core CPU with this small model:
-#   500 rows    -> a sanity-test run, a few minutes
-#   2,000 rows  -> ~30-60 minutes, good for a first real pass
-#   5,000 rows  -> a few hours
-#   10,000+ rows -> likely half a day or more; only do this once
-#                   you've confirmed the pipeline works end-to-end
-#
-# Start SMALL (MAX_ROWS = 500) to sanity check, then raise it.
-MAX_ROWS = 2000
-
-# Fraction of MAX_ROWS held out for validation
-VAL_FRACTION = 0.1
-
-# Training hyperparameters (scaled down for CPU)
+# !!! IMPORTANT - READ THIS !!!
+# Your currently-deployed model is saved as a .gguf (Ollama/llama.cpp format).
+# GGUF is an inference-only quantized format - it CANNOT be loaded by
+# transformers/PEFT to continue fine-tuning it. You need the ORIGINAL
+# merged HF/safetensors folder instead - i.e. the folder that came out of
+# the "merge_and_unload()" step in your Kaggle notebook (Cell 9), before
+# it was converted to GGUF for Ollama. Copy that folder onto this server
+# and point BASE_MODEL_PATH at it.
+BASE_MODEL_PATH = "/path/to/your/merged_model"   # <-- SET THIS (NOT a .gguf file)
+
+# Path to your HBL data. CSV must have columns named exactly: urdu, roman
+CSV_PATH = "./pairs_clean.csv"
+
+# ADJUST THIS: pairs_clean.csv has ~137k rows total.
+# Set N_ROWS to an integer (e.g. 40000) to train on a random subset,
+# or set to None to use the full ~137k rows.
+N_ROWS = None   # <-- ADJUST THIS
+
+# Fraction of the (possibly subsetted) data held out for validation
+VAL_FRACTION = 0.02
+
+# Where training checkpoints + the new LoRA adapter are saved
+OUTPUT_DIR = "./hbl-stage2-finetune"
+ADAPTER_DIR = f"{OUTPUT_DIR}/final_adapter"
+
+# Where the FINAL merged model (your old model + new adapter, merged) is saved
+MERGED_MODEL_DIR = f"{OUTPUT_DIR}/merged_model_v2"
+
+MAX_SEQ_LEN = 224
+
+# Training hyperparameters (A16, ~16GB VRAM, bf16 no quantization -
+# if you hit OOM, lower PER_DEVICE_BATCH_SIZE first, e.g. to 2, and raise GRAD_ACCUM_STEPS to compensate)
+PER_DEVICE_BATCH_SIZE = 4
+GRAD_ACCUM_STEPS = 4
 NUM_EPOCHS = 1
-PER_DEVICE_BATCH_SIZE = 1     # lowered from 2 -> less RAM per step
-GRAD_ACCUM_STEPS = 16         # effective batch size = 1 * 16 = 16 (unchanged)
 LEARNING_RATE = 2e-4
-MAX_SEQ_LENGTH = 224
+WARMUP_STEPS = 200
+SAVE_STEPS = 500
+EVAL_STEPS = 500
+
+RESPONSE_TEMPLATE = "### Response:\n"
+
 
 # ============================================================
-# 1. LOAD YOUR HBL DATASET
+# STEP 1: LOAD + CLEAN pairs_clean.csv
 # ============================================================
 print("=" * 60)
-print("Loading HBL dataset...")
+print("Loading pairs_clean.csv")
 print("=" * 60)
 
-df = pd.read_csv(DATA_PATH)
-df = df.rename(columns={"urdu": "Urdu_Input", "roman": "Roman_Urdu_Target"})
-df = df.dropna(subset=["Urdu_Input", "Roman_Urdu_Target"])
+df = pd.read_csv(CSV_PATH)
+df = df.rename(columns={c: c.strip().lower() for c in df.columns})
+assert "urdu" in df.columns and "roman" in df.columns, (
+    f"Expected columns 'urdu' and 'roman', got: {list(df.columns)}"
+)
+print(f"Loaded {len(df):,} total rows from {CSV_PATH}")
 
-print(f"Total rows available in data.csv: {len(df):,}")
+# Shuffle before subsetting so the sample isn't biased by CSV ordering
+df = df.sample(frac=1, random_state=42).reset_index(drop=True)
 
-if len(df) > MAX_ROWS:
-    df = df.sample(n=MAX_ROWS, random_state=42).reset_index(drop=True)
-    print(f"Sampled down to MAX_ROWS = {MAX_ROWS:,}")
+if N_ROWS is not None:
+    df = df.iloc[:N_ROWS].reset_index(drop=True)
+    print(f"Using a subset of {len(df):,} rows (N_ROWS={N_ROWS})")
 else:
-    print(f"Using all {len(df):,} rows (fewer than MAX_ROWS limit)")
-
-# ============================================================
-# 2. CLEAN DATA
-# ============================================================
-URDU_COL = "Urdu_Input"
-ROMAN_COL = "Roman_Urdu_Target"
+    print(f"Using all {len(df):,} rows")
 
 
 def clean_urdu(s):
@@ -118,6 +121,7 @@ def clean_urdu(s):
 
 
 def clean_roman(s):
+    # Uppercased to match the casing convention of your existing model
     if not isinstance(s, str) or len(s.strip()) == 0:
         return ""
     s = unicodedata.normalize("NFC", s)
@@ -128,63 +132,79 @@ def clean_roman(s):
     return s.upper()
 
 
-def preprocess_df(df, max_urdu_len=200, max_roman_len=250):
-    df = df.copy()
-    df[URDU_COL] = df[URDU_COL].apply(clean_urdu)
-    df[ROMAN_COL] = df[ROMAN_COL].apply(clean_roman)
-    df = df[(df[URDU_COL].str.len() > 0) & (df[ROMAN_COL].str.len() > 0)]
-    df = df.drop_duplicates(subset=[URDU_COL, ROMAN_COL])
-    df = df[df[URDU_COL].str.len() <= max_urdu_len]
-    df = df[df[ROMAN_COL].str.len() <= max_roman_len]
-    return df.reset_index(drop=True)
+df["urdu"] = df["urdu"].apply(clean_urdu)
+df["roman"] = df["roman"].apply(clean_roman)
+df = df[(df["urdu"].str.len() > 0) & (df["roman"].str.len() > 0)]
+df = df.drop_duplicates(subset=["urdu", "roman"])
+df = df[df["urdu"].str.len() <= 200]
+df = df[df["roman"].str.len() <= 250]
+df = df.reset_index(drop=True)
+print(f"After cleaning: {len(df):,} rows")
 
+train_df, val_df = train_test_split(df, test_size=VAL_FRACTION, random_state=42)
+print(f"Train: {len(train_df):,} | Val: {len(val_df):,}")
 
-df = preprocess_df(df)
-print(f"After cleaning: {len(df):,} usable rows")
 
-# Train/val split
-val_size = max(1, int(len(df) * VAL_FRACTION))
-val_df = df.iloc[:val_size].reset_index(drop=True)
-train_df = df.iloc[val_size:].reset_index(drop=True)
+# ============================================================
+# STEP 2: LOAD TOKENIZER (local only, from your existing model)
+# ============================================================
+print("\nLoading tokenizer from local model folder...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH, local_files_only=True)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+
+def format_example(urdu, roman):
+    return (
+        "### Instruction:\n"
+        "Transliterate the following Urdu text into Roman Urdu. \n"
+        "Output ONLY the Roman Urdu. No translation. No explanation.\n\n"
+        "### Input:\n"
+        f"{urdu}\n\n"
+        f"{RESPONSE_TEMPLATE}{roman}{tokenizer.eos_token}"
+    )
+
+
+train_texts = [format_example(u, r) for u, r in zip(train_df["urdu"], train_df["roman"])]
+val_texts = [format_example(u, r) for u, r in zip(val_df["urdu"], val_df["roman"])]
+
+train_dataset = Dataset.from_dict({"text": train_texts})
+val_dataset = Dataset.from_dict({"text": val_texts})
+
+print("\nSample formatted example:")
+print(train_texts[0])
 
-print(f"Train: {len(train_df):,} | Val: {len(val_df):,}")
-print("\nSample cleaned row:")
-print(f"  Urdu : {train_df.iloc[0][URDU_COL]}")
-print(f"  Roman: {train_df.iloc[0][ROMAN_COL]}")
 
 # ============================================================
-# 3. LOAD MODEL + TOKENIZER (LOCAL, OFFLINE)
+# STEP 3: LOAD YOUR EXISTING FINE-TUNED MODEL (4-bit, for training)
 # ============================================================
 print("\n" + "=" * 60)
-print("Loading local model and tokenizer...")
+print("Loading base model (your existing fine-tuned model) in bf16")
 print("=" * 60)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-tokenizer.pad_token = tokenizer.eos_token
-tokenizer.padding_side = "right"
+gc.collect()
+torch.cuda.empty_cache()
 
+# No bitsandbytes / quantization - the model fits fine in bf16 on 16GB VRAM
+# since only the small LoRA adapter needs gradients (base weights are frozen).
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_PATH,
-    torch_dtype=torch.bfloat16,  # CPU-safe, ~half the RAM of float32
+    BASE_MODEL_PATH,
+    dtype=torch.bfloat16,
+    device_map={"": 0},
     trust_remote_code=True,
+    local_files_only=True,
 )
 model.config.use_cache = False
-model.gradient_checkpointing_enable()  # trades compute for much lower RAM use
+model.gradient_checkpointing_enable()
+model.enable_input_require_grads()  # required so gradients flow into the LoRA adapter
+# with the base model frozen + gradient checkpointing enabled
 
-# ============================================================
-# 4. LoRA CONFIG (no quantization - CPU doesn't support bitsandbytes)
-# ============================================================
 lora_config = LoraConfig(
-    r=16,                 # lowered from 32 -> lighter for CPU training
-    lora_alpha=32,
+    r=32,
+    lora_alpha=64,
     target_modules=[
-        "q_proj",
-        "k_proj",
-        "v_proj",
-        "o_proj",
-        "gate_proj",
-        "up_proj",
-        "down_proj",
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
     ],
     lora_dropout=0.05,
     bias="none",
@@ -192,39 +212,18 @@ lora_config = LoraConfig(
 )
 
 model = get_peft_model(model, lora_config)
-print("LoRA adapter attached.")
+print("Model loaded and LoRA adapter attached.")
 model.print_trainable_parameters()
 
+
 # ============================================================
-# 5. FORMAT DATA FOR INSTRUCTION TUNING
+# STEP 4: TOKENIZE + CUSTOM COMPLETION-ONLY COLLATOR
 # ============================================================
-def format_example(urdu, roman):
-    return f"""### Instruction:
-Transliterate the following Urdu text into Roman Urdu.
-Output ONLY the Roman Urdu. No translation. No explanation.
-
-### Input:
-{urdu}
-
-### Response:
-{roman}{tokenizer.eos_token}"""
-
-
-train_texts = [format_example(u, r) for u, r in zip(train_df[URDU_COL], train_df[ROMAN_COL])]
-val_texts = [format_example(u, r) for u, r in zip(val_df[URDU_COL], val_df[ROMAN_COL])]
-
-train_dataset = Dataset.from_dict({"text": train_texts})
-val_dataset = Dataset.from_dict({"text": val_texts})
-
-print(f"\nTrain samples: {len(train_dataset):,}")
-print(f"Val samples:   {len(val_dataset):,}")
-
-
 def tokenize_function(examples):
     return tokenizer(
         examples["text"],
         truncation=True,
-        max_length=MAX_SEQ_LENGTH,
+        max_length=MAX_SEQ_LEN,
         padding=False,
         return_tensors=None,
     )
@@ -233,21 +232,15 @@ def tokenize_function(examples):
 tokenized_train = train_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
 tokenized_val = val_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
 
-# ============================================================
-# 6. CUSTOM COLLATOR - mask prompt, train only on Roman Urdu response
-# ============================================================
-response_template = "### Response:\n"
-response_ids = tokenizer.encode(response_template, add_special_tokens=False)
+model.config.use_cache = False
+response_ids = tokenizer.encode(RESPONSE_TEMPLATE, add_special_tokens=False)
 
 
 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 everything up to and including "### Response:\n"
     window_size = len(response_ids)
     for i in range(labels.shape[0]):
         label_seq = labels[i].tolist()
@@ -264,7 +257,7 @@ def custom_completion_collator(features):
 
 
 # ============================================================
-# 7. TRAINING ARGUMENTS (CPU-safe)
+# STEP 5: TRAIN
 # ============================================================
 training_args = TrainingArguments(
     output_dir=OUTPUT_DIR,
@@ -273,21 +266,22 @@ training_args = TrainingArguments(
     gradient_accumulation_steps=GRAD_ACCUM_STEPS,
     num_train_epochs=NUM_EPOCHS,
     learning_rate=LEARNING_RATE,
-    warmup_steps=20,
-    logging_steps=5,
+    warmup_steps=WARMUP_STEPS,
+    logging_steps=10,
     eval_strategy="steps",
-    eval_steps=100,
+    eval_steps=EVAL_STEPS,
     eval_accumulation_steps=1,
-    save_strategy="no",          # no intermediate checkpoints - only saved at the end (step 9)
+    save_strategy="steps",
+    save_steps=SAVE_STEPS,
+    save_total_limit=2,
     load_best_model_at_end=False,
-    fp16=False,                  # CPU: no fp16
-    bf16=False,                  # leave off unless your CPU supports it well
+    bf16=True,
+    fp16=False,
     report_to="none",
-    dataloader_num_workers=0,    # safest default on Windows
+    dataloader_num_workers=2,
     remove_unused_columns=False,
-    optim="adamw_torch",         # CPU-safe optimizer (no paged_adamw_8bit)
-    push_to_hub=False,           # offline - nothing gets uploaded
-    use_cpu=True,
+    optim="adamw_torch",  # standard optimizer, no bitsandbytes dependency
+    push_to_hub=False,
 )
 
 trainer = Trainer(
@@ -298,72 +292,72 @@ trainer = Trainer(
     data_collator=custom_completion_collator,
 )
 
-# ============================================================
-# 8. TRAIN (resumes automatically if checkpoints exist)
-# ============================================================
 print("\n" + "=" * 60)
-print("Starting training...")
+print("Starting training")
 print("=" * 60)
-print(f"Rows: {len(train_dataset):,} train / {len(val_dataset):,} val")
-print("CPU training is slow - this may take a while depending on MAX_ROWS.")
-print("Note: no intermediate checkpoints are saved - if this run is")
-print("interrupted, training must be restarted from scratch.")
 
-trainer.train()
+checkpoints = (
+    [d for d in os.listdir(OUTPUT_DIR) if d.startswith("checkpoint-")]
+    if os.path.exists(OUTPUT_DIR) else []
+)
+if checkpoints:
+    print(f"Found {len(checkpoints)} existing checkpoints. Resuming...")
+    trainer.train(resume_from_checkpoint=True)
+else:
+    print("No checkpoints found. Starting fresh training...")
+    trainer.train()
+
+model.save_pretrained(ADAPTER_DIR)
+tokenizer.save_pretrained(ADAPTER_DIR)
+print(f"\nStage-2 LoRA adapter saved to: {ADAPTER_DIR}")
 
-# ============================================================
-# 9. SAVE LoRA ADAPTER
-# ============================================================
-adapter_path = os.path.join(OUTPUT_DIR, "final_adapter")
-model.save_pretrained(adapter_path)
-tokenizer.save_pretrained(adapter_path)
-print(f"\nTraining complete. LoRA adapter saved to: {adapter_path}")
 
 # ============================================================
-# 10. MERGE ADAPTER INTO BASE MODEL (produces a standalone model)
+# STEP 6: MERGE NEW ADAPTER INTO YOUR EXISTING MODEL
 # ============================================================
 print("\n" + "=" * 60)
-print("Merging LoRA adapter into base model...")
+print("Merging new adapter into your existing fine-tuned model")
 print("=" * 60)
 
-del model, trainer
+del trainer
+del model
 gc.collect()
+torch.cuda.empty_cache()
 
+print("Loading base model (full precision, CPU) for a clean merge...")
 base_model = AutoModelForCausalLM.from_pretrained(
-    MODEL_PATH,
-    torch_dtype=torch.bfloat16,
+    BASE_MODEL_PATH,
+    dtype=torch.bfloat16,
+    device_map="cpu",
     trust_remote_code=True,
+    local_files_only=True,
 )
 
-merged_model = PeftModel.from_pretrained(base_model, adapter_path)
+print("Merging weights...")
+merged_model = PeftModel.from_pretrained(base_model, ADAPTER_DIR)
 merged_model = merged_model.merge_and_unload()
 
-merged_save_dir = os.path.join(OUTPUT_DIR, "merged_model")
-merged_model.save_pretrained(merged_save_dir)
-tokenizer.save_pretrained(merged_save_dir)
-print(f"Merged model saved to: {merged_save_dir}")
+os.makedirs(MERGED_MODEL_DIR, exist_ok=True)
+merged_model.save_pretrained(MERGED_MODEL_DIR)
+tokenizer.save_pretrained(MERGED_MODEL_DIR)
+print(f"\nFinal merged model saved to: {MERGED_MODEL_DIR}")
+print("You can now convert this folder to GGUF for Ollama, same as before.")
+
 
 # ============================================================
-# 11. QUICK SANITY TEST
+# STEP 7 (OPTIONAL): QUICK SANITY CHECK
 # ============================================================
-print("\n" + "=" * 60)
-print("Running quick sanity test on merged model...")
-print("=" * 60)
-
-
-def transliterate(urdu_text, max_new_tokens=128):
-    prompt_text = 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_text, return_tensors="pt")
+def transliterate(model_for_test, urdu_text, max_new_tokens=128):
+    prompt_text = (
+        "### Instruction:\n"
+        "Transliterate the following Urdu text into Roman Urdu. \n"
+        "Output ONLY the Roman Urdu. No translation. No explanation.\n\n"
+        "### Input:\n"
+        f"{urdu_text}\n\n{RESPONSE_TEMPLATE}"
+    )
+    inputs = tokenizer(prompt_text, return_tensors="pt").to(model_for_test.device)
     with torch.no_grad():
-        outputs = merged_model.generate(
+        outputs = model_for_test.generate(
             **inputs,
             max_new_tokens=max_new_tokens,
             do_sample=False,
@@ -373,15 +367,15 @@ Output ONLY the Roman Urdu. No translation. No explanation.
     return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
 
 
-# Pull a couple of examples straight from your own val set for the test
-sample_rows = val_df.head(3) if len(val_df) > 0 else train_df.head(3)
-for _, row in sample_rows.iterrows():
-    result = transliterate(row[URDU_COL])
-    print(f"Urdu     : {row[URDU_COL]}")
-    print(f"Expected : {row[ROMAN_COL]}")
-    print(f"Predicted: {result}")
-    print("-" * 70)
-
-print("\nDone. Model artifacts:")
-print(f"  LoRA adapter : {adapter_path}")
-print(f"  Merged model : {merged_save_dir}")
+if __name__ == "__main__":
+    print("\n" + "=" * 60)
+    print("Running a quick sanity check on the merged model")
+    print("=" * 60)
+    test_model = merged_model.to("cuda") if torch.cuda.is_available() else merged_model
+    sample_rows = df.sample(min(5, len(df)), random_state=1)
+    for _, row in sample_rows.iterrows():
+        result = transliterate(test_model, row["urdu"])
+        print(f"Urdu  : {row['urdu']}")
+        print(f"Roman (expected) : {row['roman']}")
+        print(f"Roman (model)    : {result}")
+        print("-" * 70)