Update README.md

README.md

@@ -1,86 +1,79 @@
 ---
 license: apache-2.0
-library_name: transformers_zamba2
+datasets:
+- HuggingFaceH4/ultrachat_200k
+- BAAI/Infinity-Instruct
+- HuggingFaceH4/ultrafeedback_binarized
+- Intel/orca_dpo_pairs
+- argilla/OpenHermesPreferences
+- BramVanroy/dolly-15k-dutch
+base_model:
+- Zyphra/Zamba2-1.2B-instruct
+library_name: transformers
 ---
 
-# Model Card for Zamba2-1.2B
+# Model Card for Zamba2-1.2B-instruct-Dutch
 
-Zamba2-1.2B is a hybrid model composed of state-space ([Mamba](https://github.com/state-spaces/mamba)) and transformer blocks. It broadly follows the [Zamba architecture](https://arxiv.org/abs/2405.16712) which consists of a Mamba backbone alternating with shared transformer blocks (see diagram in [Model Details](#model-details)). Zamba2-1.2B possesses three major improvements over Zamba1:
+Zamba2-1.2B-instruct-Dutch is a Dutch language instruction-following model obtained through a two-stage fine-tuning process:
 
-1.) Mamba1 blocks have been replaced with Mamba2 blocks.
+1. First stage (Base instruction model by Zyphra):
+   - Zyphra fine-tuned Zamba2-1.2B to create Zamba2-1.2B-instruct through:
+     - SFT training on [ultrachat_200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) and [Infinity-Instruct](https://huggingface.co/datasets/BAAI/Infinity-Instruct)
+     - DPO training on [ultrafeedback_binarized](https://huggingface.co/datasets/HuggingFaceH4/ultrafeedback_binarized), [orca_dpo_pairs](https://huggingface.co/datasets/Intel/orca_dpo_pairs), and [OpenHermesPreferences](https://huggingface.co/datasets/argilla/OpenHermesPreferences)
 
-2.) We apply a LoRA projector to each shared MLP and attention block, which allows the network to specialize at each invocation of the shared transformer layer across depth. LoRA enables us to add depth-specialization for only a minimal increase in total parameter count.
+2. Second stage (Dutch language adaptation):
+   - Further fine-tuning of Zyphra's Zamba2-1.2B-instruct on the [dolly-15k-dutch](https://huggingface.co/datasets/BramVanroy/dolly-15k-dutch) dataset, specifically using the training split
 
-3.) We utilize rotary position embeddings in the shared attention layer.
-
-Zamba2-1.2B differs from our [2.7B model](https://huggingface.co/Zyphra/Zamba2-2.7B) in three ways:
-
-1.) We have added rotary position embeddings
-
-2.) A single shared transformer block (instead of two that we alternate between)
-
-3.) Added LoRA projectors to attention blocks (instead of just a LoRA on the MLP block)
-
-We found that while hybrid SSM-transformer models are perfectly capable of performing well without position embeddings, adding rotary embeddings to the shared attention block slightly improved performance. Secondly, we utilize a single attention block (instead of alternating between two independent transformer blocks) because this enables a higher flop count for the model at a given parameter budget and at smaller scales this becomes more important than the slightly faster latency.
-
-Zamba2-1.2B uses the Mistral v0.1 tokenizer and was pre-trained on 3T tokens of text and code data sourced from open web-datasets, including [Zyda](https://arxiv.org/abs/2406.01981). Subsequently, in a second phase, Zamba2-1.2B was annealed on a mixture of 100B high-quality tokens.
-
-Note: this is a temporary HuggingFace implementation of Zamba2-1.2B. It may not yet be fully compatible with all frameworks and tools intended to interface with HuggingFace models.
-
-A standalone Pytorch implementation of Zamba2-1.2B may be found [here](https://github.com/Zyphra/Zamba2).
+The model maintains the core hybrid architecture of Zamba2 while being optimized for Dutch language understanding and generation.
 
 ## Quick start
 
 ### Prerequisites
 
-To download Zamba2-1.2B, clone Zyphra's fork of transformers:
+To download Zamba2-1.2B-instruct-Dutch, clone Zyphra's fork of transformers:
 1. `git clone https://github.com/Zyphra/transformers_zamba2.git`
 2. `cd transformers_zamba2`
 3. Install the repository: `pip install -e .`
 4. `pip install accelerate`
 
-
-You can run the model without using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly higher latency and memory usage. 
-
-To run on CPU, please specify `use_mamba_kernels=False` when loading the model using ``AutoModelForCausalLM.from_pretrained``.
-
-
 ### Inference
+
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-1.2B")
-model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-1.2B", device_map="cuda", torch_dtype=torch.bfloat16)
+# Instantiate model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-1.2B-instruct-Dutch")
+model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-1.2B-instruct-Dutch", device_map="cuda", torch_dtype=torch.bfloat16)
 
-input_text = "What factors contributed to the fall of the Roman Empire?"
-input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
+# Format the input as a chat template
+prompt = "Wat zijn de belangrijkste oorzaken van de val van het Romeinse Rijk?"
+sample = [{'role': 'user', 'content': prompt}]
+chat_sample = tokenizer.apply_chat_template(sample, tokenize=False)
 
-outputs = model.generate(**input_ids, max_new_tokens=100)
-print(tokenizer.decode(outputs[0]))
+# Tokenize input and generate output
+input_ids = tokenizer(chat_sample, return_tensors='pt', add_special_tokens=False).to("cuda")
+outputs = model.generate(**input_ids, max_new_tokens=150, return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
+print((tokenizer.decode(outputs[0])))
 ```
 
-## Training Data
+## Training Details
+
+The model was fine-tuned using the following approach:
 
-The model is fine-tuned on the **BramVanroy/dolly-15k-dutch** dataset, specifically using the training split (`train_sft`). This dataset is not SoTA, however the goal is to demonstrate the capabilities and it fits <1024 tokens.
+1. Started with the base Zamba2-1.2B-instruct model
+2. Fine-tuned on the dolly-15k-dutch dataset using optimized learning rates
+3. Implemented memory optimization through gradient checkpointing
+4. Utilized mixed precision training (bf16)
 
-### Fine-tuning with Learning Rate Optimization
+### Fine-tuning Configuration
 
-The model includes an advanced learning rate optimization system for fine-tuning, implemented through the `LROptimizerCallback` class. This callback automatically handles learning rate optimization during training. Here's how to use it:
+The model includes an advanced learning rate optimization system for fine-tuning, implemented through the `LROptimizerCallback` class:
 
 ```python
 from transformers import AutoTokenizer, Trainer
 from lr_optimizer import setup_training, LROptimizerCallback
 
-# Method 1: Using the complete setup function
-training_setup = setup_training(
-    model_name="Zyphra/Zamba2-1.2B",
-    dataset_name="your/dataset",
-    num_trials=10
-)
-trainer = training_setup['trainer']
-
-# Method 2: Using the callback directly
 callback = LROptimizerCallback(
     num_trials=10,
     lr_range=(1e-6, 1e-4)
@@ -91,21 +84,17 @@ trainer = Trainer(
     callbacks=[callback]
 )
 
-# Start training with optimized configuration
 trainer.train()
 ```
 
-The optimization process automatically:
-- Explores learning rates between 1e-6 and 1e-4 using Bayesian optimization
-- Applies Gaussian Process Regression for precise LR selection
-- Implements memory optimization through gradient checkpointing
-- Supports both fp16 and bf16 training
-
-For detailed configuration options, see the [fine-tuning documentation](link-to-docs).
+## Model Architecture
 
-## Model Details
+Zamba2-1.2B-instruct-Dutch maintains the hybrid SSM-attention architecture of the base model:
 
-Zamba2-1.2B utilizes and extends our original Zamba hybrid SSM-attention architecture. The core Zamba architecture consists of a backbone of Mamba layers interleaved with one or more shared attention layers. This attention has shared weights to minimize the parameter cost of the model. We find that concatenating the original model embeddings to the input to this attention block improves performance, likely due to better maintenance of information across depth. The Zamba2 architecture also applies LoRA projection matrices to the shared transformer blocks to gain some additional expressivity in each block and allow each shared block to specialize slightly to its own unique position while keeping the additional parameter overhead small. 
+- Backbone of Mamba2 layers interleaved with shared attention layers
+- LoRA projection matrices for shared transformer blocks
+- Rotary position embeddings in the shared attention layer
+- Concatenated original embeddings for improved information maintenance
 
 <center>
 <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/Vay6htbnBcySR3Z6NEgwj.png" width="300" alt="Zamba architecture">
@@ -113,37 +102,29 @@ Zamba2-1.2B utilizes and extends our original Zamba hybrid SSM-attention archite
 
 ## Performance
 
-Zamba2-1.2B achieves leading and state-of-the-art performance among models of <2B parameters and is competitive with some models of significantly greater size. Moreover, due to its unique hybrid SSM architecture, Zamba2-1.2B achieves extremely low inference latency and rapid generation with a significantly smaller memory footprint than comparable transformer based models. 
-
-Zamba2-1.2B's high performance and small inference compute and memory footprint renders it an ideal generalist model for on-device applications.
-
-<center>
-<img src="https://cdn-uploads.huggingface.co/production/uploads/65bc13717c6ad1994b6619e9/7Japy8VaJzKaFEjJgtWBp.png" width="700" alt="Zamba performance">
-</center>
-
-<center>
-<img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/Viwo3-bpYLFUu7cLIUFVv.png" width="800" alt="Zamba performance">
-</center>
-
-<!--
-<center>
-<img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/JVZUvVMPIpIJy9RDyohMJ.png" width="800" alt="Zamba performance">
-</center>
--->
+The model maintains the efficient inference characteristics of the base Zamba2 architecture:
+- Low latency inference
+- Rapid generation
+- Small memory footprint
 
 Time to First Token (TTFT)             |  Output Generation
 :-------------------------:|:-------------------------:
 ![image/png](https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/5lpWDLdtPPVAk8COJq7gZ.png)  |  ![image/png](https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/V2tS6eCOGbpKybEoZmOB7.png)
 
-
-And memory overhead
-
+Memory overhead:
 <center>
 <img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/m0YUmAmiVnRg6l9m10CEt.png" width="400" alt="Zamba inference and memory cost">
 </center>
 
+## Limitations
+
+- The model is primarily focused on Dutch language understanding and generation
+- Performance on other languages may be limited
+- The training dataset size is relatively small compared to larger multilingual models
+- No explicit content moderation mechanisms are included
 
-## Notice
+## License
 
-Zamba2-1.2B is a pretrained base model and therefore does not have any moderation mechanism and may output toxic or otherwise harmful language. In addition, one should not expect good instruct or chat performance, as this model was not fine-tuned for instruction following or chat.
+This model is released under the Apache 2.0 license.
 
+Note: This is a temporary HuggingFace implementation. A standalone PyTorch implementation may be found at [Zamba2 GitHub repository](https://github.com/Zyphra/Zamba2).