ISTA-MLCV/llama_2_7b_ise

Llama 2 7B Instructional Segment Embedding (ISE)

This is the Llama 2 7B model augmented and fine-tuned with the Instructional Segment Embedding (ISE) modification, trained and evaluated in the paper ASIDE: Architectural Separation of Instructions and Data in Language Models.

Model Description

Instructional Segment Embedding (ISE), introduced in Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy, embeds instruction priority information directly into the model by adding learnable segment embeddings to distinguish between instruction and data tokens. We use ISE as a baseline in our paper.

Usage

To use this model, first clone and follow the installation instructions in the official ASIDE Repository.

Inside the repository, run the following code snippet (also provided here as a script) to do inference with this model.

import torch
import deepspeed
import json
import os
from huggingface_hub import login

from model_api import CustomModelHandler  # Import your custom handler
from model_api import format_prompt  # Import your prompt formatting function

# Define your instruction and data
instruction_text = "Translate to German."
data_text = "Who is Albert Einstein?"

# Model configuration
hf_token = os.environ["HUGGINGFACE_HUB_TOKEN"]
login(token=hf_token)
embedding_type = "ise"  
base_model = "meta-llama/Llama-2-7b-hf"
model_path = "Embeddings-Collab/llama_2_7b_ise_emb_SFTv110_from_base_run_35"

# Initialize the model handler
handler = CustomModelHandler(
    model_path, 
    base_model, 
    base_model, 
    model_path, 
    None,
    0, 
    embedding_type=embedding_type, 
    load_from_checkpoint=True
)

# Initialize DeepSpeed inference engine
engine = deepspeed.init_inference(
    model=handler.model,
    mp_size=torch.cuda.device_count(),  # Number of GPUs
    dtype=torch.float16,
    replace_method='auto',
    replace_with_kernel_inject=False
)
handler.model = engine.module

# Load prompt templates
with open("./data/prompt_templates.json", "r") as f:
    templates = json.load(f)

template = templates[0]  
instruction_text = format_prompt(instruction_text, template, "system")
data_text = format_prompt(data_text, template, "user")

# Generate output
output, inp = handler.call_model_api_batch([instruction_text], [data_text])
print(output)

Citation

If you use this model, please cite our paper:

@inproceedings{
  zverev2026aside,
  title={{ASIDE}}: Architectural Separation of Instructions and Data in Language Models},
  author={Egor Zverev and Evgenii Kortukov and Alexander Panfilov and Alexandra Volkova and Rush Tabesh and Sebastian Lapuschkin and Wojciech Samek and Christoph H. Lampert},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=C81TnwHiRM}
}