madan2248c/clinical-agent-transformer-llama-31-8b

Clinical Agent Transformer

This repository contains the trained artifacts for the Clinical Agent Transformer run:

Repository: madan2248c/clinical-agent-transformer-llama-31-8b
Stage: mcp
Profile: full
Base model: meta-llama/Llama-3.1-8B-Instruct

What Is Included

This upload contains the final trained package from the Modal checkpoint volume:

• PEFT adapter files under model/
• trainable_state.pt for the custom Clinical Agent Transformer modules
• training metrics and graph-ready CSV files
• training configuration

The trainable_state.pt file is important because this project modifies the LLaMA architecture with custom MLA, MoE, recurrent, QK-Norm, post-norm, and MTP components. The adapter files alone do not represent every custom trainable module.

Repository Layout

.
├── README.md
├── metrics.json
├── training_config.json
├── training_history.csv
├── losses.csv
├── loss_curve.png
├── train_validation_loss.png
├── validation_perplexity.png
├── learning_rate.png
├── trainable_state.pt
└── model/
    ├── adapter_config.json
    └── adapter_model.safetensors

How To Clone

git lfs install
git clone https://huggingface.co/madan2248c/clinical-agent-transformer-llama-31-8b
cd clinical-agent-transformer-llama-31-8b

The repository uses Git LFS because trainable_state.pt is large.

How To Use This Model Package

This repository is not a standalone vanilla Hugging Face AutoModelForCausalLM.from_pretrained(...) checkpoint. It is a Clinical Agent Transformer package with:

1. the original gated base model dependency, meta-llama/Llama-3.1-8B-Instruct;
2. PEFT/LoRA adapter weights in model/;
3. custom architecture trainable weights in trainable_state.pt;
4. project code that applies the architecture surgery before loading the trained state.

To use it correctly, clone this project codebase and load the model through the same architecture path used during training.

Install Runtime Dependencies

pip install torch transformers==4.52.4 peft accelerate bitsandbytes huggingface_hub

You must also have access to the gated LLaMA base model and be logged in:

huggingface-cli login

Example Loading Code

import torch
from huggingface_hub import hf_hub_download
from peft import PeftModel

from llama_surgery.config import AdapterConfig, SurgeryConfig, TrainingConfig
from llama_surgery.model import load_model, load_tokenizer

repo_id = "madan2248c/clinical-agent-transformer-llama-31-8b"
hf_token = "YOUR_HF_TOKEN"

config = TrainingConfig(
    model_name="meta-llama/Llama-3.1-8B-Instruct",
    hf_cache_dir="./hf_cache",
    adapter=AdapterConfig(
        enabled=True,
        adapter_name="clinical_agent_adapter",
        r=8,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=("q_proj", "k_proj", "v_proj", "o_proj"),
    ),
    surgery=SurgeryConfig(
        use_mla=True,
        use_qk_norm=True,
        use_post_norm=True,
        use_moe_layers=True,
        use_recurrent_layers=True,
        use_mtp_head=True,
    ),
)

tokenizer = load_tokenizer(config, hf_token)
model = load_model(config, hf_token)

state_path = hf_hub_download(repo_id=repo_id, filename="trainable_state.pt", token=hf_token)
state = torch.load(state_path, map_location="cpu")
model.load_state_dict(state, strict=False)

model = PeftModel.from_pretrained(model, repo_id, subfolder="model", token=hf_token)
model.eval()

prompt = "Evaluate ceftriaxone for community-acquired pneumonia in an adult with normal renal function."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.2,
        do_sample=False,
    )

print(tokenizer.decode(output[0], skip_special_tokens=True))

Important Loading Note

The architecture surgery must be applied before loading trainable_state.pt. If someone loads only the adapter folder, they will miss the custom MLA, MoE, recurrent, post-norm, QK-Norm, and MTP trained components.

Training Artifacts

The metrics and graph files can be used to inspect the training behavior:

• metrics.json: final run metrics and architecture inspection.
• training_history.csv: train loss, validation loss, validation perplexity, and learning rate per recorded step.
• train_validation_loss.png: train/validation loss curve.
• validation_perplexity.png: validation perplexity curve.
• learning_rate.png: learning-rate schedule.

Training Procedure

The model was trained in two stages:

1. Clinical literacy training on the combined clinical corpus.
2. MCP agentic tool-call training on balanced antimicrobial stewardship traces.

The MCP full run resumes from the clinical full trainable state.

Reported Metrics

{
  "run_name": "clinical-agent-transformer-mcp-full",
  "dataset_path": "/root/generated_data/amr_mcp_traces_balanced_v1.jsonl",
  "dataset_records": 1800,
  "train_records": 1620,
  "validation_records": 180,
  "text_field": "text_trace",
  "resume_trainable_state": "/checkpoints/runs/clinical-agent-transformer-clinical-full/trainable_state.pt",
  "seq_len": 1024,
  "train_steps": 1200,
  "gradient_accumulation_steps": 8,
  "learning_rate": 2e-05,
  "warmup_steps": 100,
  "eval_every": 100,
  "checkpoint_every": 0,
  "losses": [
    8.261292934417725,
    8.163441836833954,
    8.177353203296661,
    8.168348789215088,
    8.076540052890778,
    7.815936863422394,
    7.896245896816254,
    7.62277615070343,
    7.4410223960876465,
    7.2093111872673035,
    6.975111722946167,
    6.913655698299408,
    6.826074838638306,
    6.770330011844635,
    6.720040142536163,
    6.698519468307495,
    6.608599305152893,
    6.593647480010986,
    6.469860732555389,
    6.429429352283478,
    6.311609864234924,
    6.294034957885742,
    6.333190381526947,
    6.2032288908958435,
    6.074349761009216,
    6.041495740413666,
    6.140882134437561,
    5.9048484563827515,
    5.890208125114441,
    5.841541886329651,
    5.716970205307007,
    5.772250652313232,
    5.625128269195557,
    5.614135146141052,
    5.610735833644867,
    5.650689780712128,
    5.3798434138298035,
    5.436018347740173,
    5.404404997825623,
    5.440618932247162,
    5.351420223712921,
    5.34623920917511,
    5.2838539481163025,
    5.104320526123047,
    5.148950219154358,
    5.184379279613495,
    5.016061782836914,
    5.0828863978385925,
    5.025748550891876,
    4.951754093170166,
    4.833639204502106,
    4.771369397640228,
    4.717774152755737,
    4.820539176464081,
    4.857066333293915,
    4.711103618144989,
    4.688219845294952,
    4.597307562828064,
    4.551635921001434,
    4.58421665430069,
    4.543719947338104,
    4.522010385990143,
    4.627478539943695,
    4.5739825963974,
    4.445137023925781,
    4.332988440990448,
    4.372862100601196,
    4.223928719758987,
    4.305863559246063,
    4.347627282142639,
    4.413300424814224,
    4.344864457845688,
    4.242972820997238,
    4.031886100769043,
    4.015578389167786,
    4.046565592288971,
    4.039056867361069,
    3.900683283805847,
    4.09597909450531,
    3.9002828896045685,
    3.835624784231186,
    3.958367168903351,
    3.8370602130889893,
    3.79279163479805,
    3.9009659588336945,
    3.707729309797287,
    3.931636482477188,
    3.6938535571098328,
    3.8825665414333344,
    3.7894521355628967,
    3.691390246152878,
    3.6710696518421173,
    3.6318428814411163,
    3.5484158098697662,
    3.615939825773239,
    3.608551263809204,
    3.575754702091217,
    3.3789921402931213,
    3.6145372688770294,
    3.357224851846695,
    3.2021254301071167,
    3.317003160715103,
    3.2356558740139008,
    3.2960659563541412,
    3.3148950040340424,
    3.4104406535625458,
    3.1494440734386444,
    3.035324215888977,
    3.2058411836624146,
    3.1177440881729126,
    3.0983194708824158,
    3.0439885556697845,
    3.15883806347847,
    3.017989546060562,
    3.022290915250778,
    2.783576160669327,
    2.9367659389972687,
    3.0676648914813995,
    2.87437841296196,
    2.7975026667118073,
    2.908842533826828,
    2.7392400205135345,
    2.744791567325592,
    2.844683438539505,
    2.839159846305847,
    2.68056720495224,
    2.7747884690761566,
    2.796884149312973,
    2.7488940060138702,
    2.635950803756714,
    2.6342001259326935,
    2.6462216675281525,
    2.461296558380127,
    2.521627187728882,
    2.602918356657028,
    2.590325504541397,
    2.382751762866974,
    2.4930808544158936,
    2.722613275051117,
    2.3861756324768066,
    2.371595025062561,
    2.273716986179352,
    2.5188612043857574,
    2.454538345336914,
    2.192589610815048,
    2.29652301967144,
    2.218859925866127,
    2.4165745079517365,
    2.1988307386636734,
    2.379779040813446,
    2.2218600660562515,
    2.351611077785492,
    2.089180052280426,
    2.047184944152832,
    2.2281170338392258,
    2.2691055089235306,
    2.2565267384052277,
    2.252262979745865,
    2.127178728580475,
    2.1124915182590485,
    2.1496333926916122,
    2.1908611059188843,
    2.0274095833301544,
    2.1065411269664764,
    2.0067661851644516,
    2.2191324532032013,
    2.16944819688797,
    1.9189021289348602,
    1.9401672929525375,
    2.1057000309228897,
    2.0840209424495697,
    1.7573165744543076,
    2.039648413658142,
    1.9279413521289825,
    1.9530391693115234,
    2.029738172888756,
    1.852142572402954,
    2.0695592015981674,
    1.8634213209152222,
    1.8441505432128906,
    1.8793940097093582,
    1.7200284749269485,
    1.7787855565547943,
    1.9997882544994354,
    1.8214413076639175,
    1.6856354027986526,
    1.773601919412613,
    1.873384416103363,
    1.7972178608179092,
    1.7428040653467178,
    1.675266981124878,
    1.792445808649063,
    1.6473048031330109,
    1.6831112056970596,
    1.8643322885036469,
    1.83014278113842,
    1.582194447517395,
    1.6508467942476273,
    1.9354162812232971,
    1.5165023356676102,
    1.7517023533582687,
    1.6173395961523056,
    1.5287262946367264,
    1.7858741879463196,
    1.4788846671581268,
    1.5867124646902084,
    1.5561960190534592,
    1.6554233729839325,
    1.5400921255350113,
    1.491695910692215,
    1.4784211963415146,
    1.5193863958120346,
    1.5498751401901245,
    1.5506616234779358,
    1.5459262281656265,
    1.474884569644928,
    1.5531335920095444,
    1.5825363099575043,
    1.6134756356477737,
    1.6596762537956238,
    1.6337790042161942,
    1.7885078191757202,
    1.563613697886467,
    1.524568423628807,
    1.5064586102962494,
    1.5275050550699234,
    1.426622360944748,
    1.6496621668338776,
    1.4984694570302963,
    1.352784439921379,
    1.4005490243434906,
    1.5127532333135605,
    1.3917295038700104,
    1.38415905

Intended Use

This model package is intended for project demonstration, research, and evaluation of a clinical-agent transformer architecture. It is not intended for direct clinical deployment without external validation, safety testing, and expert review.