Modularstarencoder

🚀 ModularStarEncoder-1B预训练模型

ModularStarEncoder-1B是一个在The Stack v2上进行预训练的编码器。ModularStarEncoder是一个模块化的预训练编码器，拥有五个出口点，允许用户根据下游任务进行多次出口微调。我们基于StarCoder-2构建了ModularStarEncoder，在bfloat16中，将其参数规模从150亿减少到了10亿。

我们的架构包含36个隐藏层，每个隐藏层有16个注意力头和4个键值头，采用了分组查询注意力（GQA）。该模型使用了旋转位置编码（RoPE），基本周期theta = 10^-6，隐藏维度为1024，中间层大小为12288。

为了提高效率，我们用双向自注意力替换了因果自注意力层。与使用滑动窗口注意力的StarCoder-2不同，我们选择了全注意力，以确保更高的模块化程度，避免滑动窗口机制的感受野限制。此外，我们将最大输入长度扩展到了2048个标记，与之前的代码编码器（如StarEncoder）相比，能够处理更长的代码片段。

最后，我们的实现集成了FlashAttention V2，以实现更快的推理。

🚀 快速开始

模型信息

微调版本

如需查看代码到代码和文本到代码的微调版本，请访问 ModularStarEncoder-finetuned。

💻 使用示例

基础用法

from transformers import AutoModel
from transformers import AutoTokenizer

#import the model
model = AutoModel.from_pretrained("andreagurioli1995/ModularStarEncoder", trust_remote_code=True)

#import the tokenizer, the tokenizer applies LEFT padding!
tokenizer = AutoTokenizer.from_pretrained("andreagurioli1995/ModularStarEncoder")

code_snippet = "your code to embed here"

#You should follow this pattern to embed a snippet of code 
sentence =  f"{tokenizer.sep_token}{code_snippet}{tokenizer.cls_token}"

#Tokenizing your sentence
tokenized_sensence = tokenizer(sentence, return_tensors="pt",truncation=True, max_length=2048)

#Embedding the tokenized sentence
embedded_sentence = model(**tokenized_sensence)

运行上述代码后，你将得到六个输出元素：

• last_hidden_state：模型最后一个隐藏层的表示；
• hidden_states：模型所有隐藏层的原始表示，未经过池化、归一化和投影；
• loss：如果提供了真实标签，则为损失值（在推理时为None）；
• prediction_logits：掩码语言模型头部的预测分数；
• seq_relationship_scores：上下文损失的预测分数（如果你想要有意义的分数，可以使用分隔符标记连接多个样本）；
• attentions：编码器的注意力分数。

🔧 技术细节

训练信息

我们以399万个标记的批量大小对ModularStarEncoder进行了245000个训练步骤的预训练，处理了1万亿个标记。预训练和微调是在512个NVIDIA Ampere（64GB）GPU上使用Leonardo超级计算机进行的，需要450000个GPU工作小时。

超参数

📄 许可证

该模型遵循BigCode OpenRAIL-M v1许可协议。你可以在此处找到完整协议。

📚 引用

@article{gurioli2025modeltrainallhierarchical,
      title={One Model to Train them All: Hierarchical Self-Distillation for Enhanced Early Layer Embeddings}, 
      author={Andrea Gurioli and Federico Pennino and João Monteiro and Maurizio Gabbrielli},
      year={2025},
      eprint={2503.03008},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2503.03008}, 
}