Views
No views yet
device_map="auto")pip install -r requirements.txt1import os
2import torch
3from huggingface_hub import snapshot_download
4from transformers import AutoTokenizer
5from models.alignx import build_alignx_model
6from inference import generate # includes built-in safety filter
7
8# Download model files from HuggingFace
9local_dir = snapshot_download("GautamKashyap/AlignX")
10
11# Load AlignX model with MoCaE
12model = build_alignx_model(
13 base_model_name_or_path="meta-llama/Llama-2-7b-hf",
14 finetuned_paths={
15 ax: os.path.join(local_dir, f"lora_{ax}")
16 for ax in ("helpful", "harmless", "honest")
17 },
18 task_matrix_paths={
19 ax: os.path.join(local_dir, "task_vectors", f"T_{ax}.pt")
20 for ax in ("helpful", "harmless", "honest")
21 },
22 load_in_4bit=True,
23 device_map="auto",
24)
25model.load_mocae(os.path.join(local_dir, "mocae_finetuning_plus_mocae"))
26model.eval()
27
28tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
29
30# Generate response (safety filter applied automatically)
31response = generate(model, tokenizer, "Explain the concept of neural networks.")
32print(response)1prompts = [
2 "What are three tips for better sleep?",
3 "How does photosynthesis work?",
4 "Is the Great Wall of China visible from space?",
5]
6for p in prompts:
7 print(f"Q: {p}")
8 print(f"A: {generate(model, tokenizer, p)}\n")python inference.py --prompt "Explain the concept of neural networks."python inference.py --interactive