amartya16/Agastya-RAG-2B

Agastya Decision RAG 2B

Agastya Decision RAG is a fine-tuned Qwen3.5-2B checkpoint with an Unsloth Clef decision head. It answers typed questions about a supplied state with probabilities. The intended application is to make routing and selection decisions inside a table-of-contents (ToC) guided retrieval-augmented generation (RAG) system: for example, deciding which chapter to inspect next or whether the available evidence warrants another retrieval step.

Learn more at Andra Labs and explore the Agastya RAG repository, where this model is used.

The base model architecture is Qwen3.5 (Qwen3_5ForConditionalGeneration). The separate joint_head.safetensors and joint_head_config.json hold the trained decision head. The tokenizer, chat template, processor files, model weights, and weight index are included in this repository. No separate architecture file is required by Transformers; config.json supplies the architecture metadata. For typed decision inference, load the checkpoint with Unsloth FastDecisionModel, which loads the decision head as well as the backbone. A plain Transformers load does not expose the typed decision API.

Use

Hosted API with Hugging Face Inference Endpoints

This repository includes handler.py for a custom Hugging Face Inference Endpoint. Upload the entire folder, then deploy the repository using the Inference Toolkit and a GPU instance with sufficient memory. The endpoint uses the Unsloth decision head to return typed answers and probabilities. Dedicated endpoints incur charges. Deployment and GPU compatibility have not yet been tested for this checkpoint. See the custom handler documentation.

Once the endpoint is running, call its URL with a Hugging Face access token authorized for the endpoint:

import os
import requests

response = requests.post(
    os.environ["HF_ENDPOINT_URL"],
    headers={"Authorization": f"Bearer {os.environ['HF_TOKEN']}"},
    json={"inputs": {
        "state": "Question: When is an appeal due? ToC: 1. Scope; 2. Filing deadlines; 3. Appeals.",
        "questions": {"next_section": {
            "type": "choice",
            "instructions": "Select the best section to read next.",
            "criteria": {"deadlines": "Section 2: filing deadlines", "appeals": "Section 3: appeals"},
        }},
    }},
    timeout=120,
)
response.raise_for_status()
print(response.json())

Local inference

Install a compatible GPU build of PyTorch and Unsloth, then run the included example:

pip install 'unsloth==2026.11.2'
python example_decision.py .

After upload, replace . with your Hugging Face model ID, for example python example_decision.py amartya16/Agastya-RAG-2B. You can also call the uploaded model directly from Python, without manually downloading files:

from unsloth import FastDecisionModel

repo_id = "amartya16/Agastya-RAG-2B"
model, tokenizer = FastDecisionModel.from_pretrained(
    model_name=repo_id,
    max_seq_length=2048,
    load_in_4bit=True,
)
FastDecisionModel.for_inference(model)

answers = FastDecisionModel.predict(
    model=model,
    tokenizer=tokenizer,
    state="Question: When is an appeal due? ToC: 1. Scope; 2. Filing deadlines; 3. Appeals.",
    questions={
        "next_section": {
            "type": "choice",
            "instructions": "Select the best section to read next.",
            "criteria": {
                "deadlines": "Section 2: filing deadlines",
                "appeals": "Section 3: appeals",
            },
        }
    },
)
print(answers)

from_pretrained downloads and caches the checkpoint through Hugging Face Hub. The example uses a small ToC routing state and prints the resulting answer and probabilities. Its sample output is not a reported benchmark result. The state and questions schema accepts choice, score, and noul (boolean) questions; see the typed-decisions dataset for task context.

The ToC retrieval pipeline must supply candidate section names and descriptions, fetch the selected sections, and handle citations and final answer generation. Those pipeline steps are outside this checkpoint. This model was fine-tuned on typed decisions, not on a ToC-specific training set; ToC routing quality has not been established by the evaluation below.

Training and evaluation

The supplied unsloth_decision_config.json records a QLoRA fine-tune on LocalLLaMA/typed-decisions with a soft cross-entropy objective, 34 steps, and 2 epochs. It records 150 held-out decisions, 66.0% accuracy and 0.0948 expected calibration error (ECE) for the fine-tuned model, versus 37.33% accuracy and 0.0967 ECE for the base model. These are checkpoint training-record figures; the exact held-out split, seed, and evaluation script are not included here, so they should not be read as a reproducible public leaderboard score or a ToC RAG benchmark.

Architecture and files

File	Purpose
config.json	Qwen3.5-2B backbone architecture and model configuration
model.safetensors-00001-of-00001.safetensors, model.safetensors.index.json	Merged backbone weights and index
joint_head.safetensors, joint_head_config.json	Clef typed-decision head weights and architecture
unsloth_decision_config.json	Decision layout, calibration, and training record
tokenizer.json, tokenizer_config.json, vocab.json, merges.txt, special_tokens_map.json, added_tokens.json, chat_template.jinja	Tokenizer and chat formatting
processor_config.json, preprocessor_config.json, video_preprocessor_config.json	Multimodal processor configuration inherited from Qwen3.5

The backbone config specifies 24 language layers, hidden size 2048, and a Qwen3.5 vision component. The decision-head config specifies width 512, four decoder layers, and two evidence routing layers. The saved model file is approximately 4.55 GB; the decision head is approximately 61 MB.

Limitations

Probabilities are calibrated only to the training record's held-out typed decisions, not to every downstream document collection. Test section routing and stopping decisions on your own ToC and document set before relying on them. This checkpoint does not itself index documents, retrieve passages, cite sources, or guarantee correct answers. The 262,144-token maximum in the backbone configuration is an architectural setting, not a demonstrated decision-task context length; the decision config records max_len: 4096.

Provenance and license

Base weights: unsloth/Qwen3.5-2B, derived from Qwen/Qwen3.5-2B. Training data: LocalLLaMA/typed-decisions. The included LICENSE is Apache-2.0.