astromindinc/axia-qwen-xray-7b

Axia — DeepSeek-Qwen-Xray-7B

Fine-tuned multi-modal X-ray reasoning model from the Axia project. Base model is `deepseek-ai/DeepSeek-R1-Distill-Qwen-7B`; the fine-tune adds:

• —A custom `XrayProcessor` module (~7.5 M params) that encodes a raw Chandra photon event list (time + energy pairs, plus pre-computed spectrum/PSD/metadata side branches) into a 4096-d vector matching the Qwen backbone's hidden size.
• —A LoRA adapter (r = 8) over Qwen's attention + MLP projections.
• —A special <xray> token in the vocab whose embedding is replaced at forward time by the XrayProcessor output.

Trained jointly on ~50 000 sources from the Chandra Source Catalog 2.1. Companion dataset: `astromindinc/axia-csc-corpus`.

Repository layout

This repo bundles three artifacts you can mix and match:

full/
    Everything you need to run the full LLM (Event Analyst) — PEFT
    adapter, tokenizer, and the PCA/UMAP bundle, all in one folder
    matching the on-disk layout the production server expects.

projector/
    Just the XrayProcessor weights + the constants needed to produce
    pca_64d vectors from raw event lists. ~130 MB total when combined
    with pca_umap/. Designed to run on CPU without the LLM.

pca_umap/
    The fitted Normalizer + PCA(64) + UMAP(2) pipeline used to project
    the 4096-d XrayProcessor output to the 64-d Atlas Vector Search
    index format and the 2-d UMAP scatter.

Quick start — full LLM (needs a GPU)

from huggingface_hub import snapshot_download
snapshot_download("astromindinc/axia-qwen-xray-7b", allow_patterns=["full/*", "full/**/*"], local_dir="./axia-full")

# Same loader the production server uses:
import torch
from transformers import AutoTokenizer
from peft import PeftModel
# (Plus the Qwen2_5_XrayConfig / Qwen2_5_XrayForConditionalGeneration classes
#  from `model/server/base_functions.py` in the axia repo.)

See `model/server/main.py` in the axia repo for the full loading procedure.

Quick start — CPU-only projector

For everything except the Event Analyst agent (neighbour search, custom source uploads, dE-dt embedding lookup), you only need the projector.

from huggingface_hub import snapshot_download
snapshot_download(
    "astromindinc/axia-qwen-xray-7b",
    allow_patterns=["projector/*", "pca_umap/*"],
    local_dir="./axia-projector",
)
# 130 MB total. Runs on any laptop. See model/projector/ in the axia repo
# for the FastAPI server that exposes /project against these files.

Provenance

• —Trained on ~50 000 Chandra Source Catalog 2.1 sources, joint LoRA
• —XrayProcessor optimisation, BFloat16 precision, 4× A100 (40 GB), 25 epochs.
• —The XrayProcessor was saved via PEFT's modules_to_save=["xray_processor"] alongside the LoRA adapter; the standalone projector/xray_processor.pt here was extracted from a PeftModel.from_pretrained(...).merge_and_unload() load via model.xray_processor.state_dict(). It is bit-for-bit identical to what the production server uses.

Files

• —full/adapter_config.json, full/adapter_model.safetensors — PEFT artifacts
• —full/tokenizer.json, full/tokenizer_config.json, full/special_tokens_map.json, full/chat_template.jinja — tokenizer with the added <xray> token
• —full/pca_umap/xray_umap_pca_bundle.joblib — bundled copy of the PCA+UMAP, kept here for one-stop reproducibility
• —projector/xray_processor.pt — standalone XrayProcessor state_dict
• —projector/projector_config.json — XrayProcessor __init__ args
• —projector/global_stats.json — normalisation constants (logE_mean/std, spec edges, PSD frequencies)
• —pca_umap/xray_umap_pca_bundle.joblib — standalone PCA+UMAP

License

MIT. The base Qwen weights and the CSC source data carry their own licenses; see the base model card and the CSC homepage respectively.

Citation

@misc{axia2026,
  title  = {Axia: a multi-agent decoder for Chandra X-ray sources},
  author = {AstroMind Authors},
  year   = {2026},
  note   = {TBA — paper in preparation}
}