""" ================================================================================ 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)