rtg-llama-lora

RTG-Llama (Rail Track Geometry — Llama 3.1-8B LoRA)

A LoRA adapter on top of Llama 3.1-8B-Instruct, fine-tuned to write five-section defect-investigation reports for rail track-geometry segments governed by 49 CFR Part 213, Class 3. Pairs with a separate VAE detector and SHAP attribution module in the parent project; this repository contains only the language-model adapter.

Part of the RTG framework (VAE detection + SHAP + RAG + RTG-Llama report generation). See the project repository at Abelar-Ai/rail_llm.

Training summary

Loading

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

BASE = "meta-llama/Llama-3.1-8B-Instruct"   # or NousResearch/Meta-Llama-3.1-8B-Instruct
ADAPTER = "BenGyi/rtg-llama-lora"

tokenizer = AutoTokenizer.from_pretrained(BASE)
base = AutoModelForCausalLM.from_pretrained(
    BASE,
    torch_dtype=torch.float32,          # see note on dtype below
    device_map="auto",
    low_cpu_mem_usage=True,
)
model = PeftModel.from_pretrained(base, ADAPTER)
model.eval()

Prompt format

The adapter expects a structured case payload as input. Minimal example:

prompt = (
    "You are a rail track-geometry maintenance reviewer. Produce a structured "
    "defect report with exactly the following sections: "
    "1. Justification, 2. Explainability, 3. Technical Explanation, "
    "4. Maintenance Recommendation, 5. References.\n\n"
    "Case payload:\n"
    "  VAE reconstruction error: 23.85  (threshold τ = 1.767)\n"
    "  Verdict: DEFECTIVE\n"
    "  Raw values (inches):\n"
    "    LProf62=0.310, RProf62=0.468, LAlign62=0.921, RAlign62=0.665,\n"
    "    Gage=56.479, Crosslevel=5.220\n"
    "  FRA Class 3 check: Crosslevel exceeds ±1.75 in limit.\n"
    "  Top SHAP drivers: Crosslevel (+17.64), LAlign62 (+2.29), "
    "RProf62 (+1.19).\n"
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=900, do_sample=False)
print(tokenizer.decode(out[0], skip_special_tokens=True))

For best results, ground the prompt in retrieved 49 CFR §213 passages (see the RAG setup in the parent repository).

Dtype note

On Blackwell GPUs (sm_120, e.g. RTX 5090) the model must be loaded in fp32: bf16/fp16 attention overflows on this architecture and produces NaN logits. On Hopper/Ada (A100, H100, RTX 4090) bf16 works and is faster.

License

Inherits Llama 3.1 Community License from the base model. You must agree to the Llama 3.1 license to use this adapter.

Citation

If you use this adapter, please cite the paper accompanying the project (see the GitHub repository for the current reference).