---
library_name: peft
base_model: meta-llama/Llama-3.1-8B-Instruct
tags:
- llama
- clinical-ai
- antimicrobial-stewardship
- tool-calling
- peft
- modal
---

# Clinical Agent Transformer

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

```text
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

```text
.
├── 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

```bash
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

```bash
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:

```bash
huggingface-cli login
```

### Example Loading Code

```python
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

```json
{
  "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.