akshit-fyb/embeddinggemma-2

Hindi/English -> Santali (Ol Chiki) fine-tune. Run on free Colab/Kaggle GPU.

pip install -q transformers datasets peft sacrebleu accelerate ctranslate2 sentencepiece

Needs a free HF token (flores_plus is gated: accept terms on its HF page, then huggingface-cli login).

import os, json, torch, sacrebleu from datasets import load_dataset, Dataset from transformers import (AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments, DataCollatorForSeq2Seq) from peft import LoraConfig, get_peft_model

BASE = "facebook/nllb-200-distilled-600M" # try facebook/nllb-200-1.3B if your GPU has >=16GB TGT = "sat_Olck" SRCS = ["eng_Latn", "hin_Deva"] EXTRA = "extra_data.jsonl" # optional: lines like {"src":"...","tgt":"...","src_lang":"hin_Deva"} DEV = "cuda" if torch.cuda.is_available() else "cpu"

tok = AutoTokenizer.from_pretrained(BASE, tgt_lang=TGT)

def flores(split): def get(code): ds = load_dataset("openlanguagedata/flores_plus", code, split=split) return {r["id"]: r["text"] for r in ds} sat = get(TGT) rows = [] for code in SRCS: src = get(code) rows += [{"src": src[i], "tgt": sat[i], "src_lang": code} for i in sat if i in src] return rows

train_rows = flores("dev") if os.path.exists(EXTRA): train_rows += [json.loads(l) for l in open(EXTRA, encoding="utf-8")] test_rows = flores("devtest") print("train:", len(train_rows), "test:", len(test_rows))

def chrf(model, rows, n=200): model.eval(); hyps, refs = [], [] for r in rows[:n]: tok.src_lang = r["src_lang"] x = tok(r["src"], return_tensors="pt", truncation=True, max_length=200).to(DEV) with torch.no_grad(): out = model.generate(**x, forced_bos_token_id=tok.convert_tokens_to_ids(TGT), num_beams=4, max_new_tokens=200) hyps.append(tok.decode(out[0], skip_special_tokens=True)); refs.append(r["tgt"]) return sacrebleu.corpus_chrf(hyps, [refs], word_order=2).score

model = AutoModelForSeq2SeqLM.from_pretrained(BASE).to(DEV) print("BASELINE chrF++:", chrf(model, test_rows))

def enc(r): tok.src_lang = r["src_lang"] return tok(r["src"], text_target=r["tgt"], max_length=128, truncation=True)

ds = Dataset.from_list(train_rows).map(enc, remove_columns=["src", "tgt", "src_lang"])

model = get_peft_model(model, LoraConfig(r=32, lora_alpha=64, lora_dropout=0.05, target_modules=["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"], task_type="SEQ_2_SEQ_LM")) model.print_trainable_parameters()

args = Seq2SeqTrainingArguments(output_dir="ckpt", per_device_train_batch_size=16, gradient_accumulation_steps=2, learning_rate=2e-4, num_train_epochs=6, warmup_ratio=0.05, lr_scheduler_type="cosine", fp16=torch.cuda.is_available(), logging_steps=20, save_strategy="no", report_to="none") Seq2SeqTrainer(model=model, args=args, train_dataset=ds, data_collator=DataCollatorForSeq2Seq(tok, model=model)).train()

print("TUNED chrF++:", chrf(model, test_rows)) # keep tuned only if this beats baseline model = model.merge_and_unload() model.save_pretrained("santali_merged"); tok.save_pretrained("santali_merged") os.system("ct2-transformers-converter --model santali_merged --output_dir santali_ct2 " "--quantization int8 --force") print("Done -> download santali_merged/ and santali_ct2/")