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}, 
}