AlignX

A fine-tuned LLAMA-2-7B model with Mixture of Calibrated Alignment Experts (MoCaE) for improved helpfulness, harmlessness, and honesty.

Requirements

• Python 3.9+
• GPU with 16GB+ VRAM (or 2× 8GB with device_map="auto")
• Access to meta-llama/Llama-2-7b-hf

Installation

pip install -r requirements.txt

Usage

Python API (recommended)

import os
import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from models.alignx import build_alignx_model
from inference import generate  # includes built-in safety filter

# Download model files from HuggingFace
local_dir = snapshot_download("GautamKashyap/AlignX")

# Load AlignX model with MoCaE
model = build_alignx_model(
    base_model_name_or_path="meta-llama/Llama-2-7b-hf",
    finetuned_paths={
        ax: os.path.join(local_dir, f"lora_{ax}")
        for ax in ("helpful", "harmless", "honest")
    },
    task_matrix_paths={
        ax: os.path.join(local_dir, "task_vectors", f"T_{ax}.pt")
        for ax in ("helpful", "harmless", "honest")
    },
    load_in_4bit=True,
    device_map="auto",
)
model.load_mocae(os.path.join(local_dir, "mocae_finetuning_plus_mocae"))
model.eval()

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# Generate response (safety filter applied automatically)
response = generate(model, tokenizer, "Explain the concept of neural networks.")
print(response)

Multiple prompts

prompts = [
    "What are three tips for better sleep?",
    "How does photosynthesis work?",
    "Is the Great Wall of China visible from space?",
]
for p in prompts:
    print(f"Q: {p}")
    print(f"A: {generate(model, tokenizer, p)}\n")

Command line — single prompt

python inference.py --prompt "Explain the concept of neural networks."

Command line — interactive mode

python inference.py --interactive

License

Apache 2.0. Please also respect the license of the base model.