Views
No views yet
| Property | Value |
|---|---|
| Base model | Qwen/Qwen3-8B |
| Architecture | AutoModelForTokenClassification (num_labels=1) |
| Training domain | Mathematics and code generation (multi-domain training) |
| Output | Per-token hallucination probability (sigmoid of logits) |
1from transformers import AutoTokenizer, AutoModelForTokenClassification
2import torch
3
4model_id = "mr233/TokenHD-8B-Mix"
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForTokenClassification.from_pretrained(model_id, num_labels=1)
7model.eval()
8
9problem = "What is the capital of France?"
10response = "The capital of France is London."
11
12messages = [
13 {"role": "user", "content": problem},
14 {"role": "assistant", "content": response},
15]
16input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)[:-2]
17input_tensor = torch.tensor(input_ids).unsqueeze(0)
18
19with torch.no_grad():
20 logits = model(input_ids=input_tensor).logits # shape: (1, seq_len, 1)
21
22# scores for response tokens only
23response_ids = tokenizer.encode(response, add_special_tokens=False)
24scores = torch.sigmoid(logits.squeeze(-1).squeeze(0))[-len(response_ids):]
25# scores[i] is the hallucination probability for the i-th response token1from datasets import load_dataset
2from transformers import AutoTokenizer, AutoModelForTokenClassification
3import torch
4import numpy as np
5
6def hard_f1(y_true, y_pred):
7 if max(y_true) == 0:
8 y_true, y_pred = 1 - y_true, 1 - y_pred
9 tp = np.sum((y_pred == 1) & (y_true == 1))
10 fp = np.sum((y_pred == 1) & (y_true == 0))
11 fn = np.sum((y_pred == 0) & (y_true == 1))
12 precision = tp / (tp + fp + 1e-7)
13 recall = tp / (tp + fn + 1e-7)
14 f1 = 2 * precision * recall / (precision + recall + 1e-7)
15 return precision, recall, f1
16
17model_id = "mr233/TokenHD-8B-Mix"
18tokenizer = AutoTokenizer.from_pretrained(model_id)
19model = AutoModelForTokenClassification.from_pretrained(
20 model_id, num_labels=1, torch_dtype=torch.bfloat16, device_map="auto"
21)
22model.eval()
23
24benchmarks = [
25 "tokenhd_eval_math_500",
26 "tokenhd_eval_math_aime",
27 "tokenhd_eval_math_gpqa",
28 "tokenhd_eval_math_fin_qa",
29 "tokenhd_eval_math_olym",
30 "tokenhd_eval_math_olym_phy",
31 "tokenhd_eval_code_codeelo",
32 "tokenhd_eval_code_live_code_lite",
33]
34
35for bench in benchmarks:
36 dataset = load_dataset("mr233/TokenHD-eval-data",
37 data_files=f"{bench}.jsonl", split="train")
38 f1_incor, f1_cor = [], []
39 for item in dataset:
40 token_weights_gt = np.array(item["token_weights"], dtype=np.float32)
41 gt_hard = (token_weights_gt > 0.5).astype(np.float32)
42
43 messages = [{"role": "user", "content": item["problem"]},
44 {"role": "assistant", "content": item["raw_answer"]}]
45 input_ids = tokenizer.apply_chat_template(
46 messages, tokenize=True, add_generation_prompt=False)[:-2]
47 input_tensor = torch.tensor(input_ids, device=model.device).unsqueeze(0)
48
49 with torch.no_grad():
50 logits = model(input_ids=input_tensor).logits
51 scores = torch.sigmoid(logits.squeeze(-1).squeeze(0))[-len(gt_hard):]
52 pred_hard = (scores.float().cpu().numpy() > 0.5).astype(np.float32)
53
54 _, _, f1 = hard_f1(gt_hard, pred_hard)
55 if item["correctness"] == -1:
56 f1_incor.append(f1)
57 else:
58 f1_cor.append(f1)
59
60 s_incor = np.mean(f1_incor) * 100 if f1_incor else float("nan")
61 s_cor = np.mean(f1_cor) * 100 if f1_cor else float("nan")
62 print(f"{bench:<44s} S_incor={s_incor:.2f} S_cor={s_cor:.2f}")