Add a collected sample to the storage

Event stream format: https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format

Simulate group-wise quantization.

Convert collected data to a HuggingFace Dataset and save as Dataset and Parquet files

Convert collected data into a HuggingFace Dataset and output: 1) A single large Parquet file 2) Multiple size-limited Parquet shards

Create default configuration

Create OpenAI-style logprobs.

Extract the LAST \\boxed{...} or \\fbox{...} from model output.

Finalize the collection process

Convert all collected samples to final storage formats

Generate text and extract hidden states Returns: Tuple of (generated_text, hidden_states_array)

Get the generate_stream function for inference.

Initialize the pipeline - load models and datasets

Parse a username:password file for gradio authorization.

Preprocess data by tokenizing and tracking first human positions.

处理隐藏状态的核心流程

对一个list中的所有hidden states进行MHA处理

Prompt for output from a role.

Rotates half the hidden dims of the input.

Run the main data collection process

Resample hidden states according to CovGauss definition in paper: - use_mean = False -> CovGauss-0μ: H_new ~ N(0, Σ̂) - use_mean = True

Convert configuration to dictionary

Convert sample to dictionary format

Write a shard of samples to Parquet format

返回： - 'hf_export'：目录下含 hidden_mha_state.pt（你自定义模块的权重单独导出） - 'trainer_ckpt'：目录下 pytorch_model.bin 里包含 hidden_mha（wrapper 的完整 state_dic

如果 tokenizer 缺少某些特殊 token，就添加并 resize 模型 embedding。

Conservatively estimate the byte size of one sample for shard size control

从 mp_rank_00_model_states.pt 的内容中提取 input_projector 的权重。 支持两种常见结构： 1) 顶层就有 key "input_projector"，且对应的是子 state_dict 2) 顶层有 "module"/"m

初始化MHA层权重

Initialize MHA layer weights

初始化输出投影层

优先从本地目录加载；不行则尝试从 config._name_or_path 回退；再不行抛错。

记录训练统计信息