Datasetpublic

GwendalTsang/mistral-7b-hidden-states-tpu-verification

Mistral-7B TPU hidden-state extractor verification Verification-only artifact — not a reproduction of the paper's AUROC. This dataset contains last-token embeddings and hidden states extracted from 1,000 flattened CoQA validation question/reference-answer pairs with mistralai/Mistral-7B-Instruct-v0.3. It verifies that the memory-bounded TPU v5e-1 extraction path runs successfully. It does not contain the Mistral best_answer generations or hallucination labels required to… See the full description on the dataset page: https://huggingface.co/datasets/GwendalTsang/mistral-7b-hidden-states-tpu-verification.

sourceHugging Faceupdated 2mo agoView on Hugging Face

0likes243downloads

Dataset Card

Mistral-7B TPU hidden-state extractor verification

Verification-only artifact — not a reproduction of the paper's AUROC.

This dataset contains last-token embeddings and hidden states extracted from 1,000 flattened CoQA validation question/reference-answer pairs with mistralai/Mistral-7B-Instruct-v0.3. It verifies that the memory-bounded TPU v5e-1 extraction path runs successfully. It does not contain the Mistral best_answer generations or hallucination labels required to reproduce Table 2 of Automatic Layer Selection for Hallucination Detection.

The paper authors' public repository does not contain its prepared_data records. Do not treat this artifact as evidence for or against the paper's FEPoID or hidden-state-probing AUROC claims.

Tensor layout

Each states/validation/shard-*.safetensors file contains BF16 tensors:

Key	Shape	Meaning
`embedding`	`[N, 4096]`	Last-position input-token embedding
`hidden_states`	`[N, 32, 4096]`	Last-position state for every Mistral layer

Layers 0–30 are post-transformer-block states. Layer 31 is post-final-RMSNorm, matching the tensors selected as outputs.hidden_states[1:] by the released CUDA code. metadata-validation.jsonl maps each source index to a shard and offset. inputs-validation.jsonl contains the exact prompt-answer text and its SHA-256 digest.

Memory-bounded TPU method

—Loads MistralModel, omitting the unused language-model head.
—Uses BF16, use_cache=False, XLA_NO_SPECIAL_SCALARS=1, and PyTorch/XLA JIT.
—Never enables output_hidden_states=True.
—Hooks only [:, -1, :] after each layer and performs one host transfer per batch.
—Uses static sequence buckets of 128, 256, 512, 1,024, and 2,048 tokens.
—Saves resumable SafeTensors shards rather than accumulating the dataset in RAM.

The 1,000-record extraction completed in 215.8 seconds after model placement, using the 256-, 512-, and 1,024-token buckets without truncating any record.

Loading

python

import json
from huggingface_hub import hf_hub_download
from safetensors import safe_open

repo_id = "GwendalTsang/mistral-7b-hidden-states-tpu-verification"
metadata_path = hf_hub_download(
    repo_id, "metadata-validation.jsonl", repo_type="dataset"
)
row = json.loads(open(metadata_path, encoding="utf-8").readline())
shard_path = hf_hub_download(repo_id, row["shard"], repo_type="dataset")
with safe_open(shard_path, framework="pt", device="cpu") as shard:
    embedding = shard.get_tensor("embedding")[row["offset"]]
    hidden_states = shard.get_tensor("hidden_states")[row["offset"]]

See `scripts/extract_mistral_hidden_states_tpu.py` for reproduction and for processing paper-compatible JSONL records containing context, question, best_answer, and label. The extractor supports both --answer-view full and the authors' rule-based --answer-view first_sentence.

Provenance

—Model revision: c170c708c41dac9275d15a8fff4eca08d52bab71
—CoQA revision: 0d9e9952f1ef6e5415492d3d84b5873259137e3c
—Paper code revision: def3cb6d262c11d252e6c5a6b7e04375b94b54db
—Paper: https://arxiv.org/abs/2605.26366
—Released code: https://github.com/DesoloYw/Automatic-Layer-Selection-for-Hallucination-Detection