Support for encoder-only models in the F2LLM framework

[harrison-huan-liu]

Summary of Changes

This PR introduces support for encoder-only models in the F2LLM framework and includes testing for compatibility with different flash attention versions. The changes enable the system to work with both encoder-only models (like BERT) and decoder-only models (like Qwen), automatically detecting the model type and applying appropriate configurations.

Related Issue

#10

Key Features Added

1. Encoder-Only Model Support: Added functionality to detect and properly handle encoder-only models (BERT, ELECTRA, MPNet) vs. decoder-only models.

2. Flexible Embedding Extraction: Implemented different embedding extraction methods based on model type:

   • CLS token pooling for encoder-only models
   • Mean pooling for encoder-only models with no pooler
   • Last token extraction for decoder-only models

3. Flash Attention Compatibility: Updated flash attention version requirements and implementation selection logic.

4. Tokenizer Script Update: Modified data tokenization process to support both BERT and Qwen tokenizers with a new --tokenizer parameter.

Technical Changes

• Added detect_model_type() function to automatically identify model architecture
• Implemented three embedding extraction methods: extract_cls_embeddings(), extract_mean_pooling_embeddings(), and extract_last_token_embeddings()
• Updated model loading logic to use appropriate attention implementation based on model type and flash attention version
• Modified tokenize_data.py to support multiple tokenizer types via command line argument
• Added BERT-specific configuration file
• Updated README with new tokenization instructions

Testing

• Tested model training with flash_attn versions 2.6.0 and 2.3.6
• Verified compatibility with both encoder-only (BERT) and decoder-only (Qwen) models
• Updated accelerate configuration for encoder-only model training

[harrison-huan-liu]

#10

[harrison-huan-liu]

#10