Aquila-135M-Intermediate

Introduction#

The Aquila-135M model is a small bilingual(Chinese and English) language model, which is trained using a two-phrase paradigm: pre-training and annealing. This model used 1.66TB bilingual tokens in Chinese and English during pre-training phrase and 100B tokens during annealing training phrase. In annealing stage, we selected 100B tokens of high-quality bilingual data and finally got our model.

The Aquila-135M-Instuct model is finetuned using Infinity Instruct.

The entire training process was conducted using FlagGems based on Triton and parallel training framework named FlagScale.

Also, we have open-sourced all intermediate checkpoints.

News#

• 2024/12/24: We have released Aquila-135M and Aquila-135M-Instruct.
• 2024/12/24: We have released all datasets and intermediate checkpoints during training. Please feel free to use these models for analysis and experimentation.

Datasets#

We have open-sourced all bilingual datasets during both pre-training and annealing phrases. Datasets composition and mix proportions are shown in the figure below. Image: datasets composition

Evaluation#

We followed the same evaluation setting of SmolLM models and evaluated models using the lighteval tool.

The parameter count excludes the embedding part and Aquila-135M and SmolLM2-135M share an identical model structure. Aquila-135M achieves comparable performance on English benchmarks, while Aquila-135M demonstrates significantly better results on Chinese benchmarks.

Among small models with a total parameter count below and around 400M, Aquila-135M maintains a leading position in processing capabilities while significantly enhancing Chinese language proficiency.

For comparison models, evaluations were conducted in a local environment, so the scores may differ slightly from those reported in papers.

How to use#

Instruct Model#

from transformers import AutoModelForCausalLM, AutoTokenizer
checkpoint = "BAAI/Aquila-135M-Instruct"

device = "cuda" # for GPU usage or "cpu" for CPU usage
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
# for multiple GPUs install accelerate and do `model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")`
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)

messages = [{"role": "user", "content": "什么是引力？"}]
input_text=tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
print(input_text)
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=500)
print(tokenizer.decode(outputs[0]))
## 引力是宇宙中的一个基本力，由多个物体相互作用而产生的。它由能量和质量组成，与引力定律密切相关。

messages = [{"role": "user", "content": "What is gravity?"}]
input_text=tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
print(input_text)
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=500)
print(tokenizer.decode(outputs[0]))
## Gravity is the force that keeps us on Earth as we orbit it. It pulls objects towards each other with a strength that depends on how far apart they are from each other, and how strong the gravitational pull is. The stronger the object's mass, the greater its gravitational pull.

Future Plan#

• We plan to further optimize the composition and proportions of the dataset.
• We plan to further explore the application of small-scale models in specific scenarios.

Citation#

If you find this useful, please cite the following work

@misc{aquila-135m,
      title={Aquila-135M: A Bilingual Small Language Model in Chinese and English}, 
      author={BAAI},
      year={},
      eprint={},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={}, 
}

Acknowledgements#

This work is supported by the National Science and Technology Major Project (No. 2022ZD0116300). 本项目受新一代人工智能国家科技重大专项（No. 2022ZD0116300）支持。