team-gemmra/gemmra

Gemmra — Pharmacovigilance LoRA Adapter for Gemma 4 31B

Gemmra is a LoRA adapter that transforms Google's Gemma 4 31B-IT into a specialized pharmacovigilance assessment system. It automates four critical drug safety tasks that typically take 30 minutes per case manually — completing them in under 10 seconds with auditable reasoning traces.

Built for the TCS & AMD AI Hackathon 2026 on AMD Instinct MI300X (192 GB HBM3).

⚠️ Research Use Only. This model is for research and educational purposes. It does not provide professional medical or regulatory advice. Do not use for clinical decision-making without expert oversight.

Key Results

Base Model Comparison

Evaluated on the same eval samples (base model used hand-crafted format prompts for fair comparison).

Model Details

• —Base Model: google/gemma-4-31b-it
• —Method: LoRA SFT (bf16, r=64) (WiSE-FT weight interpolation explored for reasoning recovery)
• —Training Hardware: AMD Instinct MI300X (192 GB HBM3)
• —Precision: bf16 (zero quantization — MI300X VRAM enables full precision)
• —Training Time: ~1.9 hours
• —VRAM Usage: 95 GB (training) / 61 GB (inference)

LoRA Configuration

WiSE-FT (Weight Interpolation Exploration)

While pure SFT (α=1.0) is the primary model deployed due to its superior accuracy across 3 out of 4 tasks and 100% format compliance, we also explored WiSE-FT as a research variant to recover reasoning depth. Scaling the LoRA adapter weights by α=0.9 blends SFT format compliance with base model reasoning depth. This recovers the base model's native clinical reasoning (providing 400+ words of structured thinking) at a small cost of ~4% composite accuracy.

θ_final = α × θ_SFT + (1 - α) × θ_base (via LoRA adapter weight scaling)

Training Data

• —Training pairs: 32,355 instruction-completion pairs
• —Eval samples: 3,560 (content-hash decontaminated, MeditronFO-inspired splitting)
• —Diversity: 93–99% unique completions via Combinatorial Diversity Engine

Data Challenges Solved

1. 1.MedDRA is proprietary — engineered PT training from BioDEX open literature
2. 2.FDA redacts doctor narratives — built structured prompts from remaining FAERS fields
3. 3.BioDEX truncation — abstracts cut at 500 chars hid ground truth from 92% of T2 data; fixing this single line gave 2.1× improvement
4. 4.Train/eval leakage — content-hash splitting ensures zero contamination

Usage

Loading the Adapter

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model (requires ~62 GB VRAM in bf16)
base_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-31b-it",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-31b-it")

# Load Gemmra LoRA adapter
model = PeftModel.from_pretrained(base_model, "Amaltrkmr/gemmra")

Running Inference

messages = [
    {"role": "system", "content": "You are a pharmacovigilance expert. Assess whether this adverse event case is SERIOUS per ICH E2A criteria (Death, Life-threatening, Hospitalization, Disability, Congenital anomaly). Think step by step, then provide your structured assessment."},
    {"role": "user", "content": """Patient: 69-year-old female
Drug: ACTEMRA (tocilizumab)
Adverse events: Cardiac arrest, Pulmonary embolism, Acute kidney injury, Haemodialysis, Platelet count decreased
Outcome: Patient did not survive"""}
]

prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=1024, temperature=0.1, do_sample=True)

response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

Expected Output:

SERIOUS: YES
Criteria met: DE (Death), LT (Life-threatening), HO (Hospitalization), DS (Disability)
Rationale: The clinical outcome meets multiple seriousness categories, confirming serious classification.

Using with Unsloth (Faster)

from unsloth import FastLanguageModel
import torch

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="google/gemma-4-31b-it",
    max_seq_length=8192,
    load_in_4bit=False,
    dtype=torch.bfloat16,
)

from peft import PeftModel
model = PeftModel.from_pretrained(model, "Amaltrkmr/gemmra")
FastLanguageModel.for_inference(model)

Four Pharmacovigilance Tasks

Training Pipeline

FAERS + BioDEX + OnSIDES
    ↓
Combinatorial Diversity Engine → 32,355 pairs
    ↓
SFT (bf16 LoRA r=64 on MI300X, ~1.9 hrs) → Primary Adapter ✅
    ↓
WiSE-FT exploration (α=0.9) → Explored reasoning variant
    ↓
GRPO validation → +0.003 composite improvement → validated SFT ceiling
    ↓
Evaluation (3,560 decontaminated samples)
    ↓
This Adapter ✅

Hardware Requirements

AMD MI300X Advantage

Training this model at bf16 precision with LoRA r=64 across all 7 linear layer types requires 95 GB VRAM. This is physically impossible on any single NVIDIA GPU (A100/H100 max at 80 GB). AMD MI300X's 192 GB HBM3 is the enabling technology — zero quantization means higher quality gradients and a better final model.

Limitations

• —MedDRA vocabulary: Trained on BioDEX-derived PTs (~5,000 terms), not the full proprietary MedDRA dictionary (80,000+ PTs). T2 accuracy will improve with dictionary augmentation.
• —Data source: FDA FAERS data has known limitations — doctor narratives are redacted, outcome codes can be inconsistent.
• —Not a medical device: Outputs require expert review before regulatory submission.
• —English only: Trained exclusively on English-language adverse event reports.

Citation

@misc{gemmra2026,
  title={Gemmra: Multi-Task Pharmacovigilance Assessment with Fine-Tuned Gemma 4 on AMD MI300X},
  author={Amal T R and Bhaskar Jha},
  year={2026},
  howpublished={TCS \& AMD AI Hackathon 2026},
  url={https://github.com/bhaskarjha-dev/gemmra}
}

Contributors

• —[Amal T R](https://huggingface.co/Amaltrkmr) — Model training, evaluation, data pipeline, WiSE-FT research
• —[Bhaskar Jha](https://huggingface.co/bhaskarjha-dev) — Architecture, data engineering, website, presentation, system design

Links

• —🌐 Website: gemmra.bhaskarjha.dev
• —💻 GitHub: bhaskarjha-dev/gemmra (upstream: amaltr/gemmra)
• —🏆 Hackathon: TCS & AMD AI Hackathon 2026 — Track: Fine-Tuning (FINETUNING_005)