| --- |
| language: |
| - en |
| license: llama2 |
| base_model: meta-llama/Llama-2-7b-chat-hf |
| tags: |
| - alignment |
| - helpfulness |
| - harmlessness |
| - honesty |
| - lora |
| --- |
| |
| # 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](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf) |
|
|
| ```bash |
| pip install -r requirements.txt |
| ``` |
|
|
| ## Usage |
|
|
| **Single prompt:** |
| ```bash |
| HF_TOKEN=your_token python inference.py --prompt "What is climate change?" |
| ``` |
|
|
| **Interactive mode:** |
| ```bash |
| HF_TOKEN=your_token python inference.py --interactive |
| ``` |
|
|
| **In Python:** |
| ```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. |
|
|