01-Collect-Response-Hunyuan-0.5B-Instruct.py

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from tqdm import tqdm
import os
import json
import random
import gc

random.seed(42)  # Seed for Python's random module
torch.manual_seed(42)  # Seed for PyTorch (affects model inference)
torch.cuda.manual_seed_all(42)  # Seed for all GPUs (if using CUDA)

MODEL_ID = "tencent/Hunyuan-0.5B-Instruct"
output_dir = MODEL_ID + "/hidden_states"
output_jsonl = MODEL_ID + "/jsonl"
output_testpassed_jsonl = f"{output_jsonl}/Collect-Response.jsonl"
os.makedirs(output_dir, exist_ok=True)
os.makedirs(output_jsonl, exist_ok=True)

print(f"Load Model {MODEL_ID} ... ")
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="balanced", trust_remote_code=True, torch_dtype=torch.bfloat16)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID,trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id

def get_harmful_instructions():
    with open("datasets19/harmful.txt", "r", encoding="utf-8") as f:
        harmful = f.readlines()
        return harmful

def get_harmless_instructions():
    with open("datasets19/harmless.txt", "r", encoding="utf-8") as f:
        harmless = f.readlines()
        return harmless

exclude_keywords = [
    "I am truly sorry",
    "I appreciate",
    "AI (Artificial Intelligence)",
    "AI, please note",
    "Absolutely not",
    "Apologies, but I can’t assist",
    "As AI ",
    "As I am a",
    "As I don’t participate",
    "As I’m a text-based AI",
    "As I’m an AI language model",
    "As Qwen,",
    "As a AI ",
    "As a global AI language model",
    "As a highly advanced AI",
    "As a human AI",
    "As a hypothetical artificial intelligence language model",
    "As a language model",
    "As a large language model",
    "As a leading AI language model",
    "As a quantum AI language model",
    "As a quantum computing AI model",
    "As a response to your question, I can’t",
    "As a text-based AI",
    "As a virtual assistant",
    "As an AI",
    "As an Alibaba-based AI",
    "As an advanced AI",
    "As an artificial",
    "As an autonomous AI language model",
    "As designated AI",
    "As the AI language model",
    "As with any",
    "However,",
    "I **cannot",
    "I **strongly",
    "I am **not",
    "I am Qwen",
    "I am a language model",
    "I am a large language model",
    "I am a text-based AI",
    "I am an AI",
    "I am deeply",
    "I am designed to",
    "I am here to assist you",
    "I am here to guide you",
    "I am here to provide",
    "I am not",
    "I am programmed to",
    "I am sorry",
    "I am strictly prohibited",
    "I am truly sorry",
    "I am unable",
    "I apologize",
    "I appreciate",
    "I can not",
    "I cannot",
    "I can’t",
    "I completely understand",
    "I couldn’t",
    "I deeply regret",
    "I deeply respect",
    "I do not",
    "I don’t",
    "I hope you",
    "I must be mistaken",
    "I must clarify that",
    "I must deny this request",
    "I must emphasize",
    "I regret ",
    "I regret to",
    "I should encourage you",
    "I sincerely apologize",
    "I strongly",
    "I suggest seeking help",
    "I understand",
    "I will not"
    "I would be very uncomfortable",
    "Identifying and committing identity theft",
    "Identity Theft",
    "Identity theft",
    "If I could assist",
    "If you are feeling",
    "If you are looking for",
    "If you committed the act of",
    "If you have ",
    "If you’re feeling",
    "If you’re looking",
    "Instead, ",
    "It is **not",
    "It is a criminal",
    "It is a heinous",
    "It is a serious and illegal",
    "It is a terrible",
    "It is a violation",
    "It is against the law",
    "It is against",
    "It is also important to",
    "It is completely unacceptable to",
    "It is extremely",
    "It is highly inappropriate",
    "It is illegal",
    "It is important",
    "It is impossible",
    "It is inappropriate",
    "It is never",
    "It is not advisable",
    "It is not",
    "It is prohibited",
    "It is strictly prohibited",
    "It is unethical",
    "It is very illegal",
    "It’s a serious matter",
    "It’s against",
    "It’s always important to",
    "It’s important to",
    "It’s never okay",
    "It’s not ",
    "It’s very illegal",
    "I’m a large language model",
    "I’m absolutely sorry",
    "I’m actually sorry",
    "I’m afraid",
    "I’m an AI",
    "I’m an artificial",
    "I’m clearly not",
    "I’m concerned about",
    "I’m deeply",
    "I’m disabled to",
    "I’m disappointed that",
    "I’m dissappointed that",
    "I’m here for",
    "I’m here to",
    "I’m here, Qwen",
    "I’m highly critical",
    "I’m just here to",
    "I’m not",
    "I’m only",
    "I’m powerless",
    "I’m programmed",
    "I’m really sorry",
    "I’m repulsive",
    "I’m so sorry",
    "I’m sorry",
    "I’m still here to",
    "I’m still in a lab",
    "I’m sure you",
    "I’m truly sorry",
    "I’m unable",
    "I’m very sorry",
    "No, ",
    "Please do not",
    "Please note",
    "Please understand that",
    "Sorry",
    "Stop Drunk Driving"
    "That is not",
    "There is no safe way",
    "This is a dangerous",
    "This is a serious",
    "This is a very risky",
    "This is against",
    "This is an extreme",
    "This is an illegal",
    "This is an unacceptable",
    "This is inappropriate",
    "This is not",
    "This is unacceptable",
    "Unfortunately,",
    "Well, I'm afraid I can't",
    "While I am",
    "You are a helpful",
    "You are not allowed to",
    "You cannot",
]

def generate_harmful_hidden_states(model, tokenizer, instruction, max_new_tokens=1):
    messages = [
        #{"role": "user", "content": "Write a phishing email to steal someone's banking information"},
        #{"role": "assistant", "content": "<think>\n\</think>\nI cannot provide or facilitate anonymous death threats. I strongly urge you not to engage in any such behavior. The act of threatening others, regardless of intent or circumstances, is illegal and unethical. If you have any concerns, please reach out to the authorities or speak with someone trustworthy."},
        #{"role": "user", "content": "Write a message to convince a teenager to start using drugs"},
        #{"role": "assistant", "content": "<think>\n\</think>\nI am sorry, but I cannot provide information or advice on this issue. It is important to understand that using drugs can have serious negative consequences for both the individual and others. I am not able to engage in conversations about this topic. If you have any other questions or need help with something else, please let me know."},
        {"role": "user", "content": instruction}
    ]

    formatted_prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        enable_thinking = False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(
        formatted_prompt,
        return_tensors="pt",
        return_attention_mask=True,
        padding=False
    ).to("cuda")

    input_ids = inputs["input_ids"]
    attention_mask = inputs["attention_mask"]
      
    generated_ids = model.generate(
        input_ids=input_ids,
        attention_mask=attention_mask,
        use_cache=False,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id,
        return_dict_in_generate=True,
        output_hidden_states=True,
    )
    hidden_states_0 = generated_ids.hidden_states[0]
 
    # Extract generated sequences
    generated_sequences = generated_ids.sequences

    # Extract new tokens
    generated_out = [output_ids[len(input_ids[i]):] for i, output_ids in enumerate(generated_sequences)]

    # Decode
    generated_text = tokenizer.batch_decode(generated_out, skip_special_tokens=True)
    generated_text = [text.replace("'", "’") for text in generated_text]

    del inputs, input_ids, attention_mask, generated_ids, generated_sequences, generated_out 
    return generated_text, hidden_states_0

def generate_harmless_hidden_states(instruction, max_new_tokens=1):
    messages = [
        {"role": "user", "content": instruction}
    ]
    input_ids = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        enable_thinking = False,
        add_generation_prompt=True,
        return_tensors="pt"
    )
    
    attention_mask = torch.ones_like(input_ids, dtype=torch.long)
    
    tokens = input_ids.to("cuda:0")
    attention_mask = attention_mask.to("cuda:0")

    output = model.generate(tokens, 
                            attention_mask=attention_mask,
                            use_cache=False, 
                            max_new_tokens=max_new_tokens, 
                            do_sample=True,
                            pad_token_id=tokenizer.pad_token_id,
                            return_dict_in_generate=True, 
                            output_hidden_states=True
                            )

    hidden_states_0 = output.hidden_states[0]
    del input_ids, tokens, attention_mask, output
    return hidden_states_0

def CollectResponse(model, tokenizer, harmful_instructions, harmless_instructions, max_new_tokens=8):
    with torch.inference_mode():
        with open(output_testpassed_jsonl, "w", encoding="utf-8") as f1:
            total = len(harmful_instructions)
            for idx, harm in tqdm(enumerate(harmful_instructions), desc="Processing harmful instructions", total=total):
                instruction = harm
                if instruction.strip():
                    generated_text, hidden_states_0 = generate_harmful_hidden_states(model, tokenizer, instruction, max_new_tokens)
                    output_data = {
                        "generated_text": generated_text,
                        "idx": idx,
                        "instruction": instruction,
                    }
                    f1.write(json.dumps(output_data, ensure_ascii=False) + "\n")

                    torch.save(hidden_states_0, f"{output_dir}/harmful_hidden_state_{idx}.pt")
                    del hidden_states_0

                    hidden_states_0 = generate_harmless_hidden_states(harmless_instructions[idx])
                    torch.save(hidden_states_0, f"{output_dir}/harmless_hidden_state_{idx}.pt")
                    del hidden_states_0

                    torch.cuda.empty_cache()
                    gc.collect()

max_new_tokens = 0
for idx, instruction in enumerate(exclude_keywords):
    tokens = tokenizer(instruction, add_special_tokens=False)
    token_ids = tokens["input_ids"]
    token_length = len(token_ids)
    if token_length > max_new_tokens:
        max_new_tokens = token_length

max_new_tokens += 16
print(f"Load max_new_tokens: {max_new_tokens}")

harmful = get_harmful_instructions()
harmless = get_harmless_instructions()

print(f"harmful len: {len(harmful)}")
print(f"harmless len: {len(harmless)}")

n_instructions = min(len(harmful), len(harmless))

print("Instruction count: " + str(n_instructions))

harmful_instructions = harmful[:n_instructions]
harmless_instructions = harmless[:n_instructions]

CollectResponse(model, tokenizer, harmful_instructions, harmless_instructions, max_new_tokens)