#!/usr/bin/env python # -*- coding: utf-8 -*- """ Reproduction script for the AlexandriaX NileChat-3B QLoRA adapters ================================================================== AlexandriaX-2026 Subtask 1 - QLoRA fine-tunes of UBC-NLP/NileChat-3B-Base (the same base model as the organisers' official baseline) for English -> dialectal Arabic dialogue translation. One script, three variants - the ablation this pair of repos exists to document: --variant noctx the prompt's history block is always "No previous turns" -> NAMAA-Space/alexandriax-nilechat-lora dev 23.54 spBLEU / 39.68 chrF++ (BEST of the three) --variant ctx the model's own previous turns are fed back as history -> not released; dev 22.87 / 39.11 --variant ctx_aux ctx + 348,787 back-translated pairs -> NAMAA-Space/alexandriax-nilechat-ctx-aux dev 22.77 / 38.71 "dev" is the official 12,250-turn AlexandriaX development set, country-macro spBLEU (sacrebleu tokenize=flores200) and chrF++ (word_order=2). Conditioning on the conversation COSTS 0.67 spBLEU here, and the back-translation augmentation does not win it back on the official split - see the model cards for the full analysis. This is the exact recipe of the original Colab notebook (`AlexandriaX_NB2_Finetune.ipynb`), extracted into a single runnable file. pip install transformers trl peft bitsandbytes datasets sacrebleu accelerate python train_nilechat_qlora.py --variant noctx --out ./runs/nilechat_noctx python train_nilechat_qlora.py --variant ctx_aux --out ./runs/nilechat_ctx_aux Hardware used originally: one A100-80GB. Author: NAMAA Community (Fatimah Emad Eldin, Omer Nacar, Khloud Al Jallad, Mona Abdelazim) License: CC-BY-NC-4.0 - check UBC-NLP/NileChat-3B-Base for the base model's terms """ import argparse import gc import glob import json import os import re from collections import defaultdict from pathlib import Path import torch from datasets import Dataset, get_dataset_config_names, get_dataset_split_names, load_dataset from peft import LoraConfig, prepare_model_for_kbit_training from transformers import (AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainerCallback) from trl import SFTConfig, SFTTrainer # -------------------------------------------------------------------------------------- # Configuration - the released adapters' exact values # -------------------------------------------------------------------------------------- MODEL_NAME = "UBC-NLP/NileChat-3B-Base" DATASET = "UBC-NLP/alexandria" AUX_DATASET = "NAMAA-Space/alexandria-backtranslated-pairs" # 348,787 synthetic pairs MAX_SEQ_LENGTH = 1024 MAX_NEW_TOKENS = 128 EPOCHS = 1 BATCH = 32 # per device (A100-80GB) GRAD_ACCUM = 2 # -> effective batch 64; 63,130 / 64 = ~987 steps for noctx/ctx LR = 2e-4 WARMUP_RATIO = 0.03 LR_SCHEDULER = "cosine" OPTIM = "paged_adamw_8bit" LOGGING_STEPS, SAVE_STEPS, SAVE_LIMIT = 20, 200, 2 SEED = 42 USE_4BIT = True LORA_R, LORA_ALPHA, LORA_DROPOUT = 16, 32, 0.05 # "all-linear" resolves on this architecture to: # q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj LORA_TARGETS = "all-linear" COMPLETION_ONLY_LOSS = True # the prompt is masked out of the loss PACKING = False DIALECT_NAME = {"EG": "Egyptian", "JO": "Jordanian", "LB": "Lebanese", "LY": "Libyan", "MA": "Moroccan", "MR": "Mauritanian", "OM": "Omani", "PS": "Palestinian", "SA": "Saudi", "SD": "Sudanese", "SY": "Syrian", "TN": "Tunisian", "YE": "Yemeni"} # -------------------------------------------------------------------------------------- # The prompt - identical to the organisers' baseline, and to the aux dataset's `prompt` field # -------------------------------------------------------------------------------------- def format_history(history): if not history: return "No previous turns (Start of conversation)." return "\n".join(f"{h.get('speaker') or 'Speaker'}: {h['sentence']}\n" f"Translation: {h['translation']}" for h in history) def build_prompt(record, turn, history): dialect = record.get("dialect") or "Arabic Dialect" domain = record.get("domain") or "Unknown Domain" participants = record.get("participants") or "Unknown Participants" direction = turn.get("direction") or "Unknown" speaker = turn.get("speaker") or "Unknown Speaker" return (f"You are an expert translator. Translate the English sentence into {dialect}.\n\n" f"### Metadata:\n- Country: {record.get('country', '')}\n- Domain: {domain}\n" f"- Participants: {participants}\n- Speaker: {speaker}\n" f"- Speaker Direction: {direction}\n\n" f"### Conversation History:\n{format_history(history)}\n\n" f"### Sentence to Translate:\n{turn.get('sentence', '').strip()}\n\n" f"### Translation:\n") def clean_generation(text): text = text.strip() for marker in ("\n###", "### Sentence to Translate:", "### Translation:"): if marker in text: text = text.split(marker, 1)[0].strip() return text # -------------------------------------------------------------------------------------- # Data # -------------------------------------------------------------------------------------- def _sorted_turns(turns): return sorted(turns or [], key=lambda t: int(t.get("turn_order", 0) or 0)) def _target_of(turn): for k in ("reference", "dialectal_translation", "translation", "target", "gold"): if turn.get(k): return turn[k] return None def build_pairs(split="train", use_context=True): """One {prompt, response} per turn. With use_context=True the history holds the GOLD previous turns at TRAINING time (that is teacher forcing, and it is correct); at EVALUATION time it must hold the model's OWN previous outputs. Mixing those up inflated an early evaluation here by roughly +2.4 spBLEU - see `translate_conversation` below for the honest loop. """ configs = set(get_dataset_config_names(DATASET)) countries = [c for c in DIALECT_NAME if c in configs and split in set(get_dataset_split_names(DATASET, c))] pairs = [] for country in countries: for row in load_dataset(DATASET, country, split=split): record = {"country": country, "dialect": row.get("dialect") or DIALECT_NAME[country], "domain": row.get("domain"), "participants": row.get("participants")} turns = _sorted_turns(row.get("turns", [])) golds = {int(t.get("turn_order", i + 1)): _target_of(t) for i, t in enumerate(_sorted_turns(row.get("dialectal_conversation", row.get("turns", []))))} history = [] for i, turn in enumerate(turns): order = int(turn.get("turn_order", i + 1)) gold = golds.get(order) if turn.get("sentence") and gold: pairs.append({"prompt": build_prompt(record, turn, history if use_context else []), "response": gold}) history.append({"speaker": turn.get("speaker", ""), "sentence": turn.get("sentence", ""), "translation": gold}) print(f"[data] {split} (context={use_context}): {len(pairs)} pairs " f"over {len(countries)} countries") return pairs def load_aux_pairs(): ds = load_dataset(AUX_DATASET, split="train") pairs = [{"prompt": r["prompt"], "response": r["response"]} for r in ds] print(f"[data] aux back-translation: {len(pairs)} pairs") return pairs def build_sft_dataset(pairs, tokenizer): """TRL's prompt/completion schema - with completion_only_loss the prompt is not trained on.""" return Dataset.from_list([{"prompt": p["prompt"], "completion": p["response"].strip() + tokenizer.eos_token} for p in pairs]) # -------------------------------------------------------------------------------------- # Training # -------------------------------------------------------------------------------------- class ProgressCallback(TrainerCallback): def __init__(self, path): super().__init__() self.path = path def on_save(self, args, state, control, **kw): try: Path(self.path).write_text(json.dumps({"global_step": state.global_step, "epoch": state.epoch})) except Exception: pass def latest_ckpt(d): cks = [p for p in glob.glob(os.path.join(d, "checkpoint-*")) if os.path.isdir(p)] return max(cks, key=lambda p: int(re.findall(r"checkpoint-(\d+)", p)[-1])) if cks else None def compute_dtype(): return (torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16) def qconfig(): if not USE_4BIT: return None return BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=compute_dtype(), bnb_4bit_use_double_quant=True) def load_tokenizer(side="right"): tok = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) if tok.pad_token is None: tok.pad_token = tok.eos_token tok.padding_side = side # right for training, left for batched generation return tok def main(): ap = argparse.ArgumentParser() ap.add_argument("--variant", choices=["noctx", "ctx", "ctx_aux"], required=True) ap.add_argument("--out", default=None) ap.add_argument("--epochs", type=int, default=EPOCHS) ap.add_argument("--batch", type=int, default=BATCH) ap.add_argument("--grad-accum", type=int, default=GRAD_ACCUM) ap.add_argument("--lr", type=float, default=LR) args = ap.parse_args() out = args.out or f"./runs/nilechat_{args.variant}" os.makedirs(out, exist_ok=True) use_context = args.variant in ("ctx", "ctx_aux") pairs = build_pairs("train", use_context=use_context) if args.variant == "ctx_aux": pairs = pairs + load_aux_pairs() # simple concatenation, no reweighting print(f"[data] mixture: {len(pairs)} examples " f"({100 * 348787 / len(pairs):.1f}% synthetic)") tok = load_tokenizer(side="right") train_ds = build_sft_dataset(pairs, tok) print(train_ds) bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported() load_kw = dict(device_map="auto", torch_dtype=compute_dtype(), quantization_config=qconfig(), trust_remote_code=True) try: model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, attn_implementation="flash_attention_2", **load_kw) except Exception: model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, **load_kw) model.config.use_cache = False if USE_4BIT: model = prepare_model_for_kbit_training(model) lora = LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, bias="none", task_type="CAUSAL_LM", target_modules=LORA_TARGETS) sft = SFTConfig( output_dir=out, num_train_epochs=args.epochs, per_device_train_batch_size=args.batch, gradient_accumulation_steps=args.grad_accum, learning_rate=args.lr, warmup_ratio=WARMUP_RATIO, lr_scheduler_type=LR_SCHEDULER, max_length=MAX_SEQ_LENGTH, completion_only_loss=COMPLETION_ONLY_LOSS, logging_steps=LOGGING_STEPS, save_steps=SAVE_STEPS, save_total_limit=SAVE_LIMIT, eval_strategy="no", bf16=bf16, fp16=not bf16, gradient_checkpointing=True, optim=OPTIM if USE_4BIT else "adamw_torch", report_to="none", packing=PACKING, seed=SEED, ) trainer = SFTTrainer(model=model, args=sft, train_dataset=train_ds, processing_class=tok, peft_config=lora, callbacks=[ProgressCallback(os.path.join(out, "progress.json"))]) ck = latest_ckpt(out) print(("resume " + ck) if ck else "fresh start") trainer.train(resume_from_checkpoint=ck) trainer.model.save_pretrained(out) tok.save_pretrained(out) print("[done] saved adapter ->", out) del trainer, model gc.collect() torch.cuda.empty_cache() # -------------------------------------------------------------------------------------- # Honest inference: history = the model's OWN previous outputs, decoded by turn position # -------------------------------------------------------------------------------------- def translate_conversation(model, tok, record, turns, use_context=True): tok.padding_side = "left" history, out = [], [] for i, turn in enumerate(_sorted_turns(turns)): prompt = build_prompt(record, turn, history if use_context else []) enc = tok([prompt], return_tensors="pt", padding=True, truncation=True, max_length=MAX_SEQ_LENGTH).to(model.device) with torch.no_grad(): gen = model.generate(**enc, max_new_tokens=MAX_NEW_TOKENS, do_sample=False) pred = clean_generation( tok.decode(gen[0][enc["input_ids"].shape[1]:], skip_special_tokens=True)) out.append(pred) if use_context: history.append({"speaker": turn.get("speaker", ""), "sentence": turn.get("sentence", ""), "translation": pred}) # own output - never the gold return out if __name__ == "__main__": main()