Views
No views yet
pip install -r requirements.txtHF_TOKEN=your_token python inference.py --prompt "What is climate change?"HF_TOKEN=your_token python inference.py --interactive1import os
2from models.base_model import load_base_model, load_tokenizer
3from models.mocae_model import TrinityXModel
4import yaml, torch
5
6with open("config.yaml") as f:
7 cfg = yaml.safe_load(f)
8
9tokenizer = load_tokenizer(cfg["backbone"], hf_token=os.environ["HF_TOKEN"])
10tokenizer.padding_side = "left"
11
12base_model = load_base_model(cfg["backbone"], precision="bfloat16",
13 device_map="auto", hf_token=os.environ["HF_TOKEN"])
14
15model = TrinityXModel.load_pretrained(
16 save_dir="trinityX_final",
17 base_model=base_model,
18 adapter_paths=["expert_helpfulness", "expert_harmlessness", "expert_honesty"],
19 mocae_config=cfg,
20)
21model.eval()
22
23prompt = "[INST] Your question here [/INST]"
24inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
25with torch.no_grad():
26 out = model.base_model.generate(**inputs, max_new_tokens=200, do_sample=False)
27print(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.