Flatten 1_raw_dataset submodules into plain tracked files
FocalLoRA, Should-It-Be-Executed-Or-Processed, and topicattack were nested git repos (with an inner FocalLoRA/data/FocalLoRA/.git as well). Drop their .git history and track the contents directly in this repo instead of as submodules/gitlinks. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
HenryChou020514
192
Codes/1_raw_dataset/FocalLoRA/code/test/GetAS.py
Normal file
192
Codes/1_raw_dataset/FocalLoRA/code/test/GetAS.py
Normal file
@ -0,0 +1,192 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Detect Important Attention Heads
|
||||
--------------------------------
|
||||
• Single-GPU: Forces model/LoRA to specified GPU; blocks non-target devices like cuda:0.
|
||||
• Multi-GPU: Exposes user-specified GPUs; uses device_map="auto" for slicing.
|
||||
"""
|
||||
|
||||
import os, json, argparse
|
||||
from pathlib import Path
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import torch, numpy as np
|
||||
from tqdm import tqdm
|
||||
from transformers import (
|
||||
AutoConfig,
|
||||
AutoTokenizer,
|
||||
AutoModelForCausalLM,
|
||||
)
|
||||
from peft import PeftModel
|
||||
|
||||
# ========================= 1. Model Loader =========================
|
||||
def load_generic_model(model_dir: str,
|
||||
device,
|
||||
device_map_cfg: Dict):
|
||||
"""
|
||||
device : torch.device('cuda:i') or cpu
|
||||
device_map_cfg : {"": i} for single-GPU or "auto" for multi-GPU
|
||||
"""
|
||||
cfg = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
||||
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
tokenizer.padding_side = "right"
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_dir,
|
||||
config=cfg,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
attn_implementation="eager",
|
||||
device_map=device_map_cfg,
|
||||
)
|
||||
return model, tokenizer
|
||||
|
||||
# ========================= 2. Score Function =========================
|
||||
def trim_and_stack(rows: List[np.ndarray]) -> np.ndarray:
|
||||
L = min(len(r) for r in rows)
|
||||
return np.stack([r[:L] for r in rows])
|
||||
|
||||
def trim_to_same(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
||||
L = min(a.shape[1], b.shape[1])
|
||||
return a[:, :L], b[:, :L]
|
||||
|
||||
def score_heads(normal: Dict[str, List[np.ndarray]],
|
||||
conflict: Dict[str, List[np.ndarray]],
|
||||
eps: float = 1e-6):
|
||||
scores = {}
|
||||
for k in normal:
|
||||
if k not in conflict:
|
||||
continue
|
||||
try:
|
||||
n = trim_and_stack(normal[k])
|
||||
c = trim_and_stack(conflict[k])
|
||||
n, c = trim_to_same(n, c)
|
||||
except Exception as e:
|
||||
print(f"⚠️ Skipped {k} (incompatible shape): {e}")
|
||||
continue
|
||||
if n.size == 0 or c.size == 0:
|
||||
continue
|
||||
|
||||
frob = np.linalg.norm(n - c, ord="fro")
|
||||
mean_shift = np.mean(np.abs(n.mean(1) - c.mean(1)))
|
||||
|
||||
def softmax(x):
|
||||
e = np.exp(x - x.max(-1, keepdims=True))
|
||||
return e / np.clip(e.sum(-1, keepdims=True), eps, None)
|
||||
|
||||
p, q = softmax(n), softmax(c)
|
||||
kl = (p * (np.log(p + eps) - np.log(q + eps))).sum() / p.shape[0]
|
||||
|
||||
scores[k] = 0.4 * frob + 0.3 * mean_shift + 0.3 * kl
|
||||
return scores
|
||||
|
||||
# ========================= 3. Extract Last-Token Attention =========================
|
||||
@torch.inference_mode()
|
||||
def extract_attention(model, tokenizer, sys_msg: str, usr_msg: str):
|
||||
msgs = [
|
||||
{"role": "system", "content": sys_msg},
|
||||
{"role": "user", "content": usr_msg},
|
||||
]
|
||||
prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
|
||||
inputs = tokenizer(prompt, return_tensors="pt")
|
||||
inputs = {k: v.to(model.device) for k, v in inputs.items()}
|
||||
|
||||
outs = model(**inputs, output_attentions=True)
|
||||
gen = model.generate(**inputs, max_new_tokens=128)
|
||||
decoded = tokenizer.decode(gen[0], skip_special_tokens=False)
|
||||
|
||||
# Extract only assistant portion
|
||||
assistant_txt = decoded.split("assistant", 1)[-1].strip() if "assistant" in decoded else decoded.strip()
|
||||
return outs.attentions, inputs["input_ids"], assistant_txt
|
||||
|
||||
# ========================= 4. Main Detection Procedure =========================
|
||||
def detect_heads(json_path: str, model, tokenizer, out_dir: str):
|
||||
with open(json_path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
grouped = defaultdict(lambda: {"normal": None, "conflict": None})
|
||||
for s in raw:
|
||||
base = s["id"].replace("_normal", "").replace("_conflict", "")
|
||||
grouped[base][s["label"]] = s
|
||||
|
||||
normal, conflict = defaultdict(list), defaultdict(list)
|
||||
responses = []
|
||||
|
||||
for _, pair in tqdm(grouped.items()):
|
||||
for lbl in ("normal", "conflict"):
|
||||
sample = pair[lbl]
|
||||
if sample is None:
|
||||
continue
|
||||
usr_msg = f"{sample['task']} {sample['user_message']}".strip() if sample["user_message"].strip() else sample["task"]
|
||||
attn, ids, output = extract_attention(model, tokenizer, sample["system_message"], usr_msg)
|
||||
|
||||
responses.append({
|
||||
"id": sample["id"], "label": lbl, "output": output
|
||||
})
|
||||
|
||||
n_layer = len(attn)
|
||||
n_head = attn[0][0].shape[0]
|
||||
last_tok = attn[0][0].shape[2] - 1
|
||||
|
||||
for L in range(n_layer):
|
||||
for H in range(n_head):
|
||||
vec = attn[L][0][H, last_tok].to(torch.float32).cpu().numpy()
|
||||
key = f"L{L}_H{H}"
|
||||
(normal if lbl == "normal" else conflict)[key].append(vec)
|
||||
|
||||
scores = score_heads(normal, conflict)
|
||||
top10 = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)[:10]
|
||||
|
||||
stem = Path(json_path).stem.replace("_instruction", "")
|
||||
tgt = Path(out_dir) / f"{stem}_outputs"
|
||||
tgt.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
out_json = tgt / "important_heads.json"
|
||||
with out_json.open("w", encoding="utf-8") as f:
|
||||
json.dump({"important_heads": [(k, float(v)) for k, v in top10],
|
||||
"responses": responses}, f, indent=2, ensure_ascii=False)
|
||||
|
||||
print(f"\n✅ Saved → {out_json}")
|
||||
print("📌 Top-10 Important Heads:")
|
||||
for h, s in top10:
|
||||
print(f" {h:8s} │ {s:8.4f}")
|
||||
|
||||
# ========================= 5. CLI Entry =========================
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--json_path", required=True)
|
||||
parser.add_argument("--model_path", required=True)
|
||||
parser.add_argument("--cuda", type=int, nargs="+", default=[0], help="GPUs to use. Example: --cuda 0 or --cuda 0 1 2")
|
||||
parser.add_argument("--output_dir", default="outputs")
|
||||
parser.add_argument("--lora_path", default="", help="Optional: LoRA adapter path")
|
||||
args = parser.parse_args()
|
||||
|
||||
# GPU setup
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join([str(i) for i in args.cuda])
|
||||
|
||||
if len(args.cuda) == 1:
|
||||
idx = args.cuda[0]
|
||||
device = torch.device(f"cuda:{idx}" if torch.cuda.is_available() else "cpu")
|
||||
device_map = {"": 0} if device.type == "cuda" else {"": "cpu"}
|
||||
else:
|
||||
device = None
|
||||
device_map = "auto"
|
||||
|
||||
print(f"🔵 Loading base model from {args.model_path} ...")
|
||||
model, tok = load_generic_model(args.model_path, device, device_map)
|
||||
|
||||
if args.lora_path:
|
||||
print(f"🟣 Loading LoRA from {args.lora_path} ...")
|
||||
model = PeftModel.from_pretrained(model, args.lora_path, device_map=device_map)
|
||||
model = model.merge_and_unload()
|
||||
print("✅ LoRA merged.")
|
||||
|
||||
detect_heads(args.json_path, model, tok, args.output_dir)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
248
Codes/1_raw_dataset/FocalLoRA/code/test/_test_mmlu.py
Normal file
248
Codes/1_raw_dataset/FocalLoRA/code/test/_test_mmlu.py
Normal file
@ -0,0 +1,248 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Utility to evaluate a base model (and optional LoRA adapter) on the MMLU benchmark.
|
||||
|
||||
The script mirrors the loading/generation settings used in `_testmodel.py` so the
|
||||
results are comparable. Pass explicit `--model_path` / `--lora_path` arguments or
|
||||
set the MODEL_PATH / LORA_PATH environment variables.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Dict, Iterable, List, Sequence, Tuple
|
||||
import tqdm
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
from peft import PeftModel
|
||||
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
|
||||
DEFAULT_MODEL = "../models/Llama-3.1-8B-Instruct/"
|
||||
DEFAULT_LORA = "../LoraAdapter/Llama-3.1-8B-Instruct_modified_0.01/batch_0"
|
||||
DEFAULT_SYSTEM_PROMPT = (
|
||||
"You are an expert tutor. Answer multiple choice questions by returning only the "
|
||||
"single letter (A, B, C, or D) for the best option. Do not add justification."
|
||||
)
|
||||
CHOICE_LETTERS = ["A", "B", "C", "D"]
|
||||
CHOICE_PATTERN = re.compile(r"\b([ABCD])\b", flags=re.IGNORECASE)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run an MMLU evaluation for a base model and optional LoRA adapter.")
|
||||
parser.add_argument("--model_path", default=os.environ.get("MODEL_PATH", DEFAULT_MODEL))
|
||||
parser.add_argument("--lora_path", default=os.environ.get("LORA_PATH", DEFAULT_LORA))
|
||||
parser.add_argument(
|
||||
"--subjects",
|
||||
type=str,
|
||||
default=os.environ.get("MMLU_SUBJECTS", "all"),
|
||||
help="Comma-separated list of MMLU subjects/configs (default: all).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
choices=["validation", "test", "train"],
|
||||
default=os.environ.get("MMLU_SPLIT", "test"),
|
||||
help="Dataset split to evaluate on.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_samples",
|
||||
type=int,
|
||||
default=int(os.environ.get("MMLU_MAX_SAMPLES", "0")),
|
||||
help="Optional cap on the number of questions per subject (0 means all).",
|
||||
)
|
||||
parser.add_argument("--system_prompt", default=DEFAULT_SYSTEM_PROMPT)
|
||||
parser.add_argument("--max_new_tokens", type=int,
|
||||
default=min(int(os.environ.get("MAX_NEW_TOKENS", "16")), 32))
|
||||
parser.add_argument("--temperature", type=float, default=0.0)
|
||||
parser.add_argument("--cuda_device", default=os.environ.get("CUDA_VISIBLE_DEVICES", "0"))
|
||||
parser.add_argument("--attn_impl", default="eager", choices=["eager", "flash_attention_2"])
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def normalize_subjects(value: str) -> List[str]:
|
||||
bits = [part.strip() for part in (value or "").split(",")]
|
||||
subjects = [part for part in bits if part]
|
||||
return subjects or ["all"]
|
||||
|
||||
|
||||
def load_tokenizer_and_model(model_path: str, attn_impl: str):
|
||||
cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=True)
|
||||
except Exception:
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False)
|
||||
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
tokenizer.padding_side = "right"
|
||||
|
||||
bnb_cfg = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
config=cfg,
|
||||
quantization_config=bnb_cfg,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
attn_implementation=attn_impl,
|
||||
)
|
||||
model.eval()
|
||||
return tokenizer, model
|
||||
|
||||
|
||||
def load_mmlu_subjects(subjects: Sequence[str], split: str, max_samples: int):
|
||||
subject_sets: List[Tuple[str, Iterable[Dict]]] = []
|
||||
for subject in subjects:
|
||||
print(f"Loading MMLU subject '{subject}' ({split} split)...")
|
||||
dataset = load_dataset("cais/mmlu", subject, split=split)
|
||||
if max_samples and max_samples > 0:
|
||||
sample_count = min(max_samples, len(dataset))
|
||||
dataset = dataset.select(range(sample_count))
|
||||
subject_sets.append((subject, dataset))
|
||||
return subject_sets
|
||||
|
||||
|
||||
def build_mmlu_prompt(tokenizer, system_prompt: str, subject: str, question: str, choices: Sequence[str]) -> str:
|
||||
choice_lines = [f"{CHOICE_LETTERS[idx]}. {choice}" for idx, choice in enumerate(choices)]
|
||||
user_message = "\n".join([
|
||||
f"Subject: {subject}",
|
||||
f"Question: {question.strip()}",
|
||||
"Choices:",
|
||||
*choice_lines,
|
||||
"Answer with only the single letter (A, B, C, or D).",
|
||||
])
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_message},
|
||||
]
|
||||
return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
|
||||
|
||||
def generate_answer(model, tokenizer, prompt: str, max_new_tokens: int, temperature: float, device: str) -> str:
|
||||
encoded = tokenizer(
|
||||
[prompt],
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
)
|
||||
torch_device = torch.device(device)
|
||||
encoded = {k: v.to(torch_device) for k, v in encoded.items()}
|
||||
with torch.inference_mode():
|
||||
outputs = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=temperature > 0,
|
||||
temperature=temperature if temperature > 0 else 1.0,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
)
|
||||
prompt_length = encoded["attention_mask"].sum(dim=1).tolist()[0]
|
||||
generated_tokens = outputs[0][prompt_length:]
|
||||
return tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
def extract_choice_letter(response: str) -> str:
|
||||
if not response:
|
||||
return ""
|
||||
match = CHOICE_PATTERN.search(response)
|
||||
if match:
|
||||
return match.group(1).upper()
|
||||
response = response.strip().upper()
|
||||
if response and response[0] in CHOICE_LETTERS:
|
||||
return response[0]
|
||||
return ""
|
||||
|
||||
|
||||
def letter_for_answer_idx(idx: int) -> str:
|
||||
if 0 <= idx < len(CHOICE_LETTERS):
|
||||
return CHOICE_LETTERS[idx]
|
||||
raise ValueError(f"Unexpected MMLU answer index: {idx}")
|
||||
|
||||
|
||||
def evaluate_model(
|
||||
model_label: str,
|
||||
model,
|
||||
tokenizer,
|
||||
subject_sets: Sequence[Tuple[str, Iterable[Dict]]],
|
||||
args,
|
||||
device: str,
|
||||
):
|
||||
total = 0
|
||||
correct = 0
|
||||
no_parse = 0
|
||||
per_subject = defaultdict(lambda: {"correct": 0, "total": 0})
|
||||
|
||||
start_time = time.time()
|
||||
for configured_subject, dataset in subject_sets:
|
||||
for idx, example in tqdm.tqdm(enumerate(dataset)):
|
||||
subject = example.get("subject", configured_subject)
|
||||
prompt = build_mmlu_prompt(tokenizer, args.system_prompt, subject, example["question"], example["choices"])
|
||||
response = generate_answer(model, tokenizer, prompt, args.max_new_tokens, args.temperature, device)
|
||||
predicted = extract_choice_letter(response)
|
||||
gold = letter_for_answer_idx(int(example["answer"]))
|
||||
|
||||
total += 1
|
||||
entry = per_subject[subject]
|
||||
entry["total"] += 1
|
||||
|
||||
if not predicted:
|
||||
no_parse += 1
|
||||
elif predicted == gold:
|
||||
correct += 1
|
||||
entry["correct"] += 1
|
||||
if args.max_samples and idx + 1 >= args.max_samples:
|
||||
break
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
accuracy = (correct / total) * 100 if total else 0.0
|
||||
|
||||
print(f"\n=== {model_label} ===")
|
||||
print(f"Questions evaluated : {total}")
|
||||
print(f"Accuracy : {accuracy:.2f}% ({correct}/{total})")
|
||||
if no_parse:
|
||||
print(f"Unparsed responses : {no_parse}")
|
||||
print(f"Elapsed time : {elapsed:.1f}s")
|
||||
print("Per-subject accuracy:")
|
||||
for subject, stats in sorted(per_subject.items()):
|
||||
subject_acc = (stats["correct"] / stats["total"]) * 100 if stats["total"] else 0.0
|
||||
print(f" {subject:30s} {stats['correct']:4d}/{stats['total']:4d} ({subject_acc:5.2f}%)")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
args.max_new_tokens = max(1, min(args.max_new_tokens, 32))
|
||||
subjects = normalize_subjects(args.subjects)
|
||||
|
||||
cuda_spec = (args.cuda_device or "").strip()
|
||||
if cuda_spec:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = cuda_spec
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
print(f"Loading base model from: {args.model_path}")
|
||||
tokenizer, model = load_tokenizer_and_model(args.model_path, attn_impl=args.attn_impl)
|
||||
|
||||
subject_sets = load_mmlu_subjects(subjects, args.split, args.max_samples)
|
||||
evaluate_model("Base model", model, tokenizer, subject_sets, args, device)
|
||||
|
||||
lora_path = (args.lora_path or "").strip()
|
||||
if lora_path:
|
||||
if not os.path.isdir(lora_path):
|
||||
raise FileNotFoundError(f"LoRA path not found: {lora_path}")
|
||||
print(f"\nApplying LoRA adapter from: {lora_path}")
|
||||
tuned_model = PeftModel.from_pretrained(model, lora_path, device_map="auto")
|
||||
evaluate_model("LoRA-tuned model", tuned_model, tokenizer, subject_sets, args, device)
|
||||
else:
|
||||
print("\n[Info] No LoRA path provided; skipping adapter evaluation.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
164
Codes/1_raw_dataset/FocalLoRA/code/test/_testmodel.py
Normal file
164
Codes/1_raw_dataset/FocalLoRA/code/test/_testmodel.py
Normal file
@ -0,0 +1,164 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Quick utility to compare a base model against a LoRA adapter on a single chat prompt.
|
||||
|
||||
The prompt is assembled from three pieces:
|
||||
* COND1 -> system message
|
||||
* COND2 -> prepended instruction in the user message (used for conflicts)
|
||||
* TASK -> the actual task/question
|
||||
|
||||
The values can come from CLI flags, environment variables, or the built-in defaults.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
|
||||
from peft import PeftModel
|
||||
|
||||
|
||||
DEFAULT_MODEL = "../models/Llama-3.1-8B-Instruct/"
|
||||
DEFAULT_LORA = "../LoraAdapter/Llama-3.1-8B-Instruct_modified_0.01/batch_0"
|
||||
DEFAULT_COND1 = "Your entire response should be in English, no other language is allowed."
|
||||
DEFAULT_COND2 = "Your entire response should be in French, no other language is allowed."
|
||||
DEFAULT_TASK = (
|
||||
"Describe the greenhouse effect and explain how human activities, "
|
||||
"such as fossil-fuel combustion, intensify this natural process."
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Compare the base model output vs. a LoRA adapter.")
|
||||
parser.add_argument("--model_path", default=os.environ.get("MODEL_PATH", DEFAULT_MODEL))
|
||||
parser.add_argument("--lora_path", default=os.environ.get("LORA_PATH", DEFAULT_LORA))
|
||||
parser.add_argument("--cond1", default=os.environ.get("COND1", DEFAULT_COND1))
|
||||
parser.add_argument("--cond2", default=os.environ.get("COND2", DEFAULT_COND2))
|
||||
parser.add_argument("--task", default=os.environ.get("TASK", DEFAULT_TASK))
|
||||
parser.add_argument("--max_new_tokens", type=int,
|
||||
default=min(int(os.environ.get("MAX_NEW_TOKENS", "512")), 512))
|
||||
parser.add_argument("--temperature", type=float, default=0.0)
|
||||
parser.add_argument("--cuda_device", default=os.environ.get("CUDA_VISIBLE_DEVICES", "0"))
|
||||
parser.add_argument("--attn_impl", default="eager", choices=["eager", "flash_attention_2"])
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_tokenizer_and_model(model_path: str, attn_impl: str):
|
||||
cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=True)
|
||||
except Exception:
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False)
|
||||
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
tokenizer.padding_side = "right"
|
||||
|
||||
bnb_cfg = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
config=cfg,
|
||||
quantization_config=bnb_cfg,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
attn_implementation=attn_impl,
|
||||
)
|
||||
model.eval()
|
||||
return tokenizer, model
|
||||
|
||||
|
||||
def build_prompt(tokenizer, cond1: str, cond2: str, task: str) -> str:
|
||||
cond1 = (cond1 or "").strip()
|
||||
cond2 = (cond2 or "").strip()
|
||||
task = (task or "").strip()
|
||||
if not cond1:
|
||||
raise ValueError("COND1/system message cannot be empty.")
|
||||
user_bits: List[str] = [x for x in (cond2, task) if x]
|
||||
user_message = " ".join(user_bits).strip()
|
||||
messages = [{"role": "system", "content": cond1}]
|
||||
if user_message:
|
||||
messages.append({"role": "user", "content": user_message})
|
||||
|
||||
return tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
|
||||
|
||||
def generate(model, tokenizer, prompt: str, max_new_tokens: int, temperature: float, device: str) -> str:
|
||||
tokenized = tokenizer(
|
||||
[prompt],
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
)
|
||||
torch_device = torch.device(device)
|
||||
tokenized = {k: v.to(torch_device) for k, v in tokenized.items()}
|
||||
with torch.no_grad():
|
||||
outputs = model.generate(
|
||||
**tokenized,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=temperature > 0,
|
||||
temperature=temperature if temperature > 0 else 1.0,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
)
|
||||
attention_mask = tokenized["attention_mask"]
|
||||
prompt_lengths = attention_mask.sum(dim=1).tolist()
|
||||
generated_tokens = outputs[0][prompt_lengths[0]:]
|
||||
return tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
def run_cases(model_label: str, model, tokenizer, prompts, max_new_tokens: int, temperature: float, device: str):
|
||||
for case_label, prompt in prompts:
|
||||
print(f"\n=== {model_label}, {case_label} ===")
|
||||
print("\nPrompt:\n")
|
||||
print(prompt)
|
||||
print("\nOutput:\n")
|
||||
response = generate(model, tokenizer, prompt, max_new_tokens, temperature, device)
|
||||
print(response)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
args.max_new_tokens = max(1, min(args.max_new_tokens, 512))
|
||||
cuda_spec = (args.cuda_device or "").strip()
|
||||
if cuda_spec:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = cuda_spec
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
print(f"Loading base model from: {args.model_path}")
|
||||
tokenizer, model = load_tokenizer_and_model(args.model_path, attn_impl=args.attn_impl)
|
||||
normal_prompt = build_prompt(tokenizer, args.cond1, "", args.task)
|
||||
conflict_prompt = build_prompt(tokenizer, args.cond1, args.cond2, args.task)
|
||||
prompt_cases = [
|
||||
("normal case", normal_prompt),
|
||||
("conflict case", conflict_prompt),
|
||||
]
|
||||
|
||||
run_cases("base model", model, tokenizer, prompt_cases, args.max_new_tokens, args.temperature, device)
|
||||
|
||||
lora_path = (args.lora_path or "").strip()
|
||||
if lora_path:
|
||||
if not os.path.isdir(lora_path):
|
||||
raise FileNotFoundError(f"LoRA path not found: {lora_path}")
|
||||
print(f"\nApplying LoRA adapter from: {lora_path}")
|
||||
lora_model = PeftModel.from_pretrained(model, lora_path, device_map="auto")
|
||||
run_cases("tuned model", lora_model, tokenizer, prompt_cases, args.max_new_tokens, args.temperature, device)
|
||||
else:
|
||||
print("\n[Info] No LoRA path provided; skipping adapter comparison.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
147
Codes/1_raw_dataset/FocalLoRA/code/test/test_sys_mask.py
Normal file
147
Codes/1_raw_dataset/FocalLoRA/code/test/test_sys_mask.py
Normal file
@ -0,0 +1,147 @@
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
from pathlib import Path
|
||||
|
||||
# -----------------------
|
||||
# Utility functions (same as main script)
|
||||
# -----------------------
|
||||
|
||||
def build_special_ids(tokenizer):
|
||||
"""Extract special token ids related to system segments."""
|
||||
sys_id = tokenizer.convert_tokens_to_ids("<|start_header_id|>")
|
||||
eot_id = tokenizer.eos_token_id
|
||||
return sys_id, eot_id
|
||||
|
||||
def make_sys_mask(input_ids: torch.Tensor, tok) -> torch.Tensor:
|
||||
"""
|
||||
Mark only the system segment based on known chat templates.
|
||||
|
||||
Supported formats:
|
||||
1. <|start_header_id|> … <|eot_id|>
|
||||
2. ChatML: <|system|> … <|end|>
|
||||
3. ChatGPT-im: <|im_start|> system … <|im_end|>
|
||||
4. LLaMA/Mistral: [INST] (system) (user) … [/INST]
|
||||
|
||||
If no format matches, returns all False mask.
|
||||
"""
|
||||
ids = input_ids
|
||||
B, L = ids.shape
|
||||
mask = torch.zeros_like(ids, dtype=torch.bool)
|
||||
|
||||
tid = tok.convert_tokens_to_ids
|
||||
start_header = tid("<|start_header_id|>")
|
||||
end_header = tid("<|end_header_id|>")
|
||||
eot = tok.eos_token_id
|
||||
sys_tok = tid("<|system|>")
|
||||
end_tok = tid("<|end|>")
|
||||
im_start = tid("<|im_start|>")
|
||||
im_end = tid("<|im_end|>")
|
||||
inst_start = tid("[INST]")
|
||||
inst_end = tid("[/INST]")
|
||||
nl_id = tid("\n")
|
||||
|
||||
for b in range(B):
|
||||
row = ids[b].tolist()
|
||||
|
||||
# Format 1: header <|start_header_id|>
|
||||
if start_header in row:
|
||||
try:
|
||||
s = row.index(end_header) + 1
|
||||
e = row.index(eot)
|
||||
if s < e:
|
||||
mask[b, s:e] = True
|
||||
continue
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Format 2: <|system|> … <|end|>
|
||||
if sys_tok in row:
|
||||
try:
|
||||
s = row.index(sys_tok) + 1
|
||||
e = row.index(end_tok, s)
|
||||
mask[b, s:e] = True
|
||||
continue
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Format 3: <|im_start|> system … <|im_end|>
|
||||
if im_start in row and im_end in row:
|
||||
for pos in [i for i, t in enumerate(row) if t == im_start]:
|
||||
if pos + 1 < L and tok.decode([row[pos + 1]]).strip() == "system":
|
||||
s = pos + 2
|
||||
try:
|
||||
e = row.index(im_end, s)
|
||||
mask[b, s:e] = True
|
||||
break
|
||||
except ValueError:
|
||||
pass
|
||||
if mask[b].any():
|
||||
continue
|
||||
|
||||
# Format 4: [INST] … [/INST]
|
||||
if inst_start in row and inst_end in row:
|
||||
ist = row.index(inst_start) + 1
|
||||
iend = row.index(inst_end)
|
||||
split = None
|
||||
blank = [(tok.decode([t]).strip() == "") for t in row[ist:iend]]
|
||||
|
||||
for idx in range(len(blank) - 1):
|
||||
if blank[idx] and blank[idx + 1]:
|
||||
split = ist + idx
|
||||
break
|
||||
if split is None:
|
||||
for idx, is_blank in enumerate(blank):
|
||||
if is_blank:
|
||||
split = ist + idx
|
||||
break
|
||||
|
||||
if split is not None and ist < split:
|
||||
mask[b, ist:split] = True
|
||||
else:
|
||||
mask[b, ist:iend] = True
|
||||
|
||||
return mask
|
||||
|
||||
def main():
|
||||
# [1] Load tokenizer (replace with your own model path)
|
||||
model_path = " "
|
||||
assert Path(model_path).exists(), f"Model path not found: {model_path}"
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
tokenizer.padding_side = "right"
|
||||
|
||||
# [2] Sample prompt for debugging
|
||||
sample = {
|
||||
"system_message": "Please always respond formally and avoid casual expressions.",
|
||||
"task": "Define the term 'machine learning'.",
|
||||
"user_message": "Make it easy to understand."
|
||||
}
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": sample["system_message"]},
|
||||
{"role": "user", "content": f"{sample['task']} {sample['user_message']}".strip()}
|
||||
]
|
||||
text_input = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
inputs = tokenizer(text_input, return_tensors="pt")
|
||||
input_ids = inputs["input_ids"]
|
||||
|
||||
# [3] Apply system mask
|
||||
sys_mask = make_sys_mask(input_ids, tokenizer)
|
||||
|
||||
# [4] Print tokens with system markers
|
||||
tokens = [tokenizer.decode([tid]) for tid in input_ids[0]]
|
||||
print("\n===== Token View with System Mask =====")
|
||||
for i, (token, is_sys) in enumerate(zip(tokens, sys_mask[0])):
|
||||
mark = "🟰" if is_sys else " "
|
||||
print(f"{i:03d} {token.strip():30s} {mark}")
|
||||
print("========================================\n")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user