TrinityX

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

Requirements

• Python 3.9+
• GPU with 16GB+ VRAM (24GB recommended)
• HuggingFace account with LLaMA-2 access

pip install -r requirements.txt

Usage

Single prompt:

HF_TOKEN=your_token python inference.py --prompt "What is climate change?"

Interactive mode:

HF_TOKEN=your_token python inference.py --interactive

In Python:

import os
from models.base_model import load_base_model, load_tokenizer
from models.mocae_model import TrinityXModel
import yaml, torch

with open("config.yaml") as f:
    cfg = yaml.safe_load(f)

tokenizer = load_tokenizer(cfg["backbone"], hf_token=os.environ["HF_TOKEN"])
tokenizer.padding_side = "left"

base_model = load_base_model(cfg["backbone"], precision="bfloat16",
                             device_map="auto", hf_token=os.environ["HF_TOKEN"])

model = TrinityXModel.load_pretrained(
    save_dir="trinityX_final",
    base_model=base_model,
    adapter_paths=["expert_helpfulness", "expert_harmlessness", "expert_honesty"],
    mocae_config=cfg,
)
model.eval()

prompt = "[INST] Your question here [/INST]"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
    out = model.base_model.generate(**inputs, max_new_tokens=200, do_sample=False)
print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Note: Always wrap your input in [INST] ... [/INST] tags for best results.