oa-rlaif-e4b-opsd-weights

Stage 1, OPSD (on-policy self-distillation): a LoRA adapter for the base model. Weights (a LoRA adapter): it needs the base model.

OA-RLAIF aligns ontologies: for each class of a source ontology it picks the corresponding class of a target ontology among retrieved candidates, or NONE. Base model: Gemma 4 E4B. Code: github.com/AGR19/OA-RLAIF (private).

The OA-RLAIF repos

Repo	Kind	What
oa-rlaif-e4b-sft-weights	weights (LoRA)	Stage 1, SFT: adapter for google/gemma-4-E4B-it
oa-rlaif-e4b-opsd-weights	weights (LoRA)	Stage 1, OPSD: adapter for google/gemma-4-E4B-it
oa-rlaif-e4b-grpo-weights	weights (LoRA)	Stage 2, GRPO: adapter for base + OPSD (merged)
oa-rlaif-e4b-sft-model-safetensors	model (safetensors, bf16)	base + SFT
oa-rlaif-e4b-opsd-grpo-model-gguf	model (GGUF Q4_K_M)	base + OPSD + GRPO
oa-rlaif-e4b-all-combined-model-safetensors	model (safetensors, bf16)	base + OPSD → GRPO + SFT
oa-rlaif-e4b-all-combined-model-gguf	model (GGUF Q4_K_M)	base + OPSD → GRPO + SFT, quantized
oa-rlaif-datasource-qa-data	data	OM datasource passages → questions and answers
oa-rlaif-opsd-sft-data	data	Stage 1 rows: the same rows train OPSD and SFT
oa-rlaif-grpo-data	data	Stage 2 prompts with their episodes; evaluation sets; splits
oa-rlaif-dral-reward-data	data	DRAL-OA vectors for the GRPO reward term r_DRAL
oa-rlaif-icd-bridge-data	data	SNOMED CT → ICD-10 → ICD-11 composed links
oa-rlaif-raw-data	data	the raw ontology folders (no PDFs), Bio-ML, parsed-ontology cache
oa-rlaif-heldout-eval-results	results	the earlier evaluation (held-out decisions)
oa-rlaif-full-eval-results	results	the full-alignment evaluation (full_eval/)

How to use

import torch
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoProcessor

base = AutoModelForImageTextToText.from_pretrained("google/gemma-4-E4B-it", dtype=torch.bfloat16)
model = PeftModel.from_pretrained(base, "gonibeed/oa-rlaif-e4b-opsd-weights")
tokenizer = AutoProcessor.from_pretrained("google/gemma-4-E4B-it").tokenizer

GRPO continues from this model: merge it (model.merge_and_unload()) before applying oa-rlaif-e4b-grpo-weights.

How the OA-RLAIF models are trained

Stage	Starts from	Data	Loss
SFT (Stage 1)	google/gemma-4-E4B-it	oa-rlaif-opsd-sft-data	cross-entropy on the reference answer y* (off-policy)
OPSD (Stage 1)	google/gemma-4-E4B-it	oa-rlaif-opsd-sft-data (the same rows)	the student samples a reply y; a frozen copy of the base model, also shown the privileged context c (the gold answer and its evidence), scores every token; loss = Σ_t KL(student ‖ teacher)
GRPO (Stage 2)	base + OPSD, merged	oa-rlaif-grpo-data, oa-rlaif-dral-reward-data	group-relative policy gradient with the reward below

All three train LoRA adapters only. LoRA r = 32, α = 64, dropout 0, on the language model's attention and MLP projections (q, k, v, o, gate, up, down); the vision and audio towers stay frozen and unused. bf16. Trained with thinking off (the chat template's enable_thinking=False). All combined = base + OPSD → GRPO (merged), then the SFT adapter's weight change added on top (task arithmetic: the SFT adapter was trained on the plain base).

Training settings

Setting	SFT	OPSD	GRPO
Learning rate, schedule	1e-4, cosine, 3% warm-up (as OPSD)	1e-4, cosine, 3% warm-up	5e-6, linear decay (trainer default), no warm-up
Batch per GPU × accumulation	8 × 4	8 × 4	8 × 2, G = 8 replies per prompt
Completion length	384 tokens	384 tokens (one sample per prompt, temperature 1.0)	256 tokens (temperature 1.0)
Other	completion-only loss	teacher = frozen base with the privileged context; top-100 logit distillation	Dr. GRPO, β = 0, rewards not scaled
Steps (from each adapter's training_args.bin)	437	539	1,068

Hardware: 4 × B200 on Vast.ai (3 GPUs train; 1 GPU serves vLLM rollouts and, in GRPO, the A4B judge).

Prompt and reply format

The model decides one source class at a time from k numbered candidates (40 for Anatomy, 20 elsewhere), with both ontologies' modelling profiles in the prompt:

[system] You are OA-RLAIF, an ontology alignment agent. Decide which candidate target entity corresponds to the
source entity. … Reply with one line starting 'THOUGHT:' that names your evidence, then exactly one bash code block.
[user]   Task: find the target entity that the source entity is equivalent to.
         Alignment: anatomy/mouse-human
         ONTOLOGY PROFILES (modelling assumptions) …
         SOURCE
         IRI: http://mouse.owl#MA_0002211
         name: left renal vein
         parents: renal vein
         CANDIDATES
         [C1] … [C40] (each with name, synonyms, parents, children, definition)

An example reply (the format the mini-SWE-agent harness executes):

THOUGHT: C7 has the same name and the parent renal vein matches.
```mswea_bash_command
echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT && om submit C7 --relation = --confidence 0.95
```

Use thinking off (enable_thinking=False), as in training, and temperature 0.

Limitations

• Trained and evaluated with thinking off. With thinking on, the model often writes its answer inside the thinking text; the code repo's oa_rlaif.eval.thinking.split recovers it.
• Decides among retrieved candidates: a correct target that retrieval misses cannot be chosen.
• The image and audio parts of Gemma 4 are present in the weights but were neither trained nor used. The Hub may list the base architecture as image-text-to-text; this model is used for text generation only.

Licence and use. Private research use only; do not make this repo public. The weights follow the Gemma terms of use. The training data contains or derives from licensed sources: the OM datasource (Springer © 2007), SNOMED CT (affiliate licence) and UMLS-derived Bio-ML references.