ruoxi_sun
58657935fc
Vendor LLMmap / llm-verify / llm-fingerprint-detector under bash/fingerprint/tools so the three fingerprint benchmarks run with only /data1/eval mounted (no /data1/xii dependency): - run.py DEFAULT_TOOLS_ROOT prefers builtin tools/, falls back to /data1/xii - exclude .git / node_modules / template backups - detector dist/ (pre-built) retained; node_modules not needed at runtime
95 lines
3.7 KiB
Python
95 lines
3.7 KiB
Python
#!/usr/bin/env python
|
|
"""Fingerprint a REAL local HF model living outside the 52-template DB (open-set demo).
|
|
|
|
Steps:
|
|
1. Load a local model (e.g. /data1/models/Qwen3-4B) directly from disk.
|
|
2. Run the 8 LLMmap fingerprinting queries against it and collect real answers.
|
|
3. Feed those answers to the LLMmap open-set predictor and print the top-K
|
|
nearest known templates.
|
|
"""
|
|
import argparse
|
|
import os
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from LLMmap.inference import load_LLMmap
|
|
|
|
|
|
def load_llm_generator(model_dir, device_map='auto', torch_dtype=torch.bfloat16, max_new_tokens=100):
|
|
tok = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
|
if tok.pad_token is None:
|
|
tok.pad_token = tok.eos_token
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_dir,
|
|
torch_dtype=torch_dtype,
|
|
device_map=device_map,
|
|
trust_remote_code=True,
|
|
)
|
|
model.eval()
|
|
|
|
def generate(query, thinking=False):
|
|
messages = [{'role': 'user', 'content': query}]
|
|
# Qwen3 chat templates accept an enable_thinking flag if the tokenizer supports it
|
|
kwargs = {}
|
|
if 'enable_thinking' in tok.chat_template:
|
|
kwargs['enable_thinking'] = thinking
|
|
prompt = tok.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True, **kwargs)
|
|
in_toks = tok(prompt, return_tensors='pt', add_special_tokens=False,
|
|
return_token_type_ids=False).to(model.device)
|
|
with torch.no_grad():
|
|
out_toks = model.generate(
|
|
**in_toks,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
pad_token_id=tok.eos_token_id,
|
|
eos_token_id=tok.eos_token_id,
|
|
)
|
|
gen = [out_toks[i, in_toks.input_ids.shape[1]:] for i in range(len(out_toks))]
|
|
return tok.batch_decode(gen, skip_special_tokens=True)[0]
|
|
|
|
return generate, tok
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description='Fingerprint a real local HF model (open-set)')
|
|
ap.add_argument('--model_dir', required=True,
|
|
help='Path to a local HuggingFace model, e.g. /data1/models/Qwen3-4B')
|
|
ap.add_argument('--model_path', default='./data/pretrained_models/default')
|
|
ap.add_argument('--save_answers', default='/tmp/real_answers.txt', help='Where to store answers')
|
|
ap.add_argument('--max_new_tokens', type=int, default=100)
|
|
ap.add_argument('--k', type=int, default=6)
|
|
ap.add_argument('--device', default='cpu', choices=['cpu', 'cuda'])
|
|
args = ap.parse_args()
|
|
|
|
model_name = os.path.basename(args.model_dir.rstrip('/'))
|
|
print(f'[1/3] Loading local model: {args.model_dir}')
|
|
generate, tok = load_llm_generator(args.model_dir, torch_dtype=torch.bfloat16,
|
|
max_new_tokens=args.max_new_tokens)
|
|
|
|
print(f'[2/3] Loading LLMmap predictor: {args.model_path}')
|
|
conf, llmmap = load_LLMmap(args.model_path, device=args.device)
|
|
|
|
print(f'[3/3] Running {len(llmmap.queries)} fingerprinting queries against {model_name}...')
|
|
answers = []
|
|
for i, q in enumerate(llmmap.queries, 1):
|
|
print(f' -- query {i}/{len(llmmap.queries)}')
|
|
try:
|
|
ans = generate(q)
|
|
except Exception as e:
|
|
ans = f'[generation error: {e}]'
|
|
answers.append(ans)
|
|
print(f' -> {ans[:160].replace(chr(10), " ")}')
|
|
|
|
with open(args.save_answers, 'w') as f:
|
|
f.write('\n'.join(answers))
|
|
print(f'\nAnswers saved to {args.save_answers}\n')
|
|
|
|
print('### LLMmap open-set prediction (target = real %s) ###' % model_name)
|
|
llmmap.print_result(llmmap(answers), k=args.k)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|