AutoMM for Named Entity Recognition in Chinese - Quick Start¶
In this tutorial, we will demonstrate how to use AutoMM for Chinese Named Entity Recognition using an e-commerce dataset extracted from one of the most popular online marketplaces, TaoBao.com. The dataset is collected and labelled by Jie et al. and the text column mainly consists of product descriptions. The following figure shows an example of Taobao product description.
Load the Data¶
We have preprocessed the dataset to make it ready-to-use with AutoMM.
import autogluon.multimodal
from autogluon.core.utils.loaders import load_pd
from autogluon.multimodal.utils import visualize_ner
train_data = load_pd.load('https://automl-mm-bench.s3.amazonaws.com/ner/taobao-ner/chinese_ner_train.csv')
dev_data = load_pd.load('https://automl-mm-bench.s3.amazonaws.com/ner/taobao-ner/chinese_ner_dev.csv')
train_data.head(5)
/home/ci/autogluon/multimodal/src/autogluon/multimodal/data/templates.py:16: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
import pkg_resources
| text_snippet | entity_annotations | |
|---|---|---|
| 0 | 雄争霸点卡/七雄争霸元宝/七雄争霸100元1000元宝直充,自动充值 | [{"entity_group": "HCCX", "start": 3, "end": 5... |
| 1 | 简约韩版粗跟艾熙百思图亲子鞋冬季百搭街头母女圆头翻边绒面厚底 | [{"entity_group": "HPPX", "start": 6, "end": 8... |
| 2 | 羚跑商务背包双肩包男士防盗多功能出差韩版休闲15.6寸电脑包皮潮 | [{"entity_group": "HPPX", "start": 0, "end": 2... |
| 3 | 热水袋防爆充电暖宝卡通毛绒萌萌可爱注水暖宫暖手宝暖水袋 | [{"entity_group": "HCCX", "start": 0, "end": 3... |
| 4 | 童装11周岁13儿童夏装男童套装2017新款10中大童15男孩12秋季5潮7 | [{"entity_group": "HCCX", "start": 0, "end": 2... |
HPPX, HCCX, XH, and MISC stand for brand, product, pattern, and Miscellaneous information (e.g., product Specification), respectively. Let’s visualize one of the examples, which is about online games top up services.
visualize_ner(train_data["text_snippet"].iloc[0], train_data["entity_annotations"].iloc[0])
雄争霸点卡 HCCX /七雄争霸 MISC 元宝 HCCX /七雄争霸 MISC 100元 MISC 1000 MISC 元宝 HCCX 直充,自动充值
Training¶
With AutoMM, the process of Chinese entity recognition is the same as English entity recognition.
All you need to do is to select a suitable foundation model checkpoint that are pretrained on Chinese or multilingual documents.
Here we use the 'hfl/chinese-lert-small' backbone for demonstration purpose.
Now, let’s create a predictor for named entity recognition by setting the problem_type to ner and specifying the label column. Afterwards, we call predictor.fit() to train the model for a few minutes.
from autogluon.multimodal import MultiModalPredictor
import uuid
label_col = "entity_annotations"
model_path = f"./tmp/{uuid.uuid4().hex}-automm_ner" # You can rename it to the model path you like
predictor = MultiModalPredictor(problem_type="ner", label=label_col, path=model_path)
predictor.fit(
train_data=train_data,
hyperparameters={'model.ner_text.checkpoint_name':'hfl/chinese-lert-small'},
time_limit=300, #second
)
=================== System Info ===================
AutoGluon Version: 1.3.2b20250716
Python Version: 3.12.10
Operating System: Linux
Platform Machine: x86_64
Platform Version: #1 SMP Wed Mar 12 14:53:59 UTC 2025
CPU Count: 8
Pytorch Version: 2.7.1+cu126
CUDA Version: 12.6
GPU Count: 1
Memory Avail: 28.44 GB / 30.95 GB (91.9%)
Disk Space Avail: 177.38 GB / 255.99 GB (69.3%)
===================================================
AutoMM starts to create your model. ✨✨✨
To track the learning progress, you can open a terminal and launch Tensorboard:
```shell
# Assume you have installed tensorboard
tensorboard --logdir /home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner
```
Seed set to 0
/home/ci/autogluon/multimodal/src/autogluon/multimodal/models/utils.py:1104: UserWarning: provided text_segment_num: 1 is smaller than hfl/chinese-lert-small's default: 2
warnings.warn(
GPU Count: 1
GPU Count to be Used: 1
Using 16bit Automatic Mixed Precision (AMP)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
| Name | Type | Params | Mode
----------------------------------------------------------------
0 | model | HFAutoModelForNER | 15.1 M | train
1 | validation_metric | MulticlassF1Score | 0 | train
2 | loss_func | CrossEntropyLoss | 0 | train
----------------------------------------------------------------
15.1 M Trainable params
0 Non-trainable params
15.1 M Total params
60.345 Total estimated model params size (MB)
232 Modules in train mode
0 Modules in eval mode
Epoch 0, global step 21: 'val_ner_token_f1' reached 0.12301 (best 0.12301), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=0-step=21.ckpt' as top 3
Epoch 0, global step 42: 'val_ner_token_f1' reached 0.59019 (best 0.59019), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=0-step=42.ckpt' as top 3
Epoch 1, global step 64: 'val_ner_token_f1' reached 0.70790 (best 0.70790), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=1-step=64.ckpt' as top 3
Epoch 1, global step 85: 'val_ner_token_f1' reached 0.73670 (best 0.73670), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=1-step=85.ckpt' as top 3
Epoch 2, global step 107: 'val_ner_token_f1' reached 0.78784 (best 0.78784), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=2-step=107.ckpt' as top 3
Epoch 2, global step 128: 'val_ner_token_f1' reached 0.81824 (best 0.81824), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=2-step=128.ckpt' as top 3
Epoch 3, global step 150: 'val_ner_token_f1' reached 0.83483 (best 0.83483), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=3-step=150.ckpt' as top 3
Epoch 3, global step 171: 'val_ner_token_f1' reached 0.81848 (best 0.83483), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=3-step=171.ckpt' as top 3
Epoch 4, global step 193: 'val_ner_token_f1' reached 0.83621 (best 0.83621), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=4-step=193.ckpt' as top 3
Epoch 4, global step 214: 'val_ner_token_f1' reached 0.85764 (best 0.85764), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=4-step=214.ckpt' as top 3
Epoch 5, global step 236: 'val_ner_token_f1' reached 0.87215 (best 0.87215), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=5-step=236.ckpt' as top 3
Epoch 5, global step 257: 'val_ner_token_f1' reached 0.84773 (best 0.87215), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=5-step=257.ckpt' as top 3
Epoch 6, global step 279: 'val_ner_token_f1' reached 0.86224 (best 0.87215), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=6-step=279.ckpt' as top 3
Epoch 6, global step 300: 'val_ner_token_f1' reached 0.88090 (best 0.88090), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=6-step=300.ckpt' as top 3
Epoch 7, global step 322: 'val_ner_token_f1' reached 0.88482 (best 0.88482), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner/epoch=7-step=322.ckpt' as top 3
Time limit reached. Elapsed time is 0:05:00. Signaling Trainer to stop.
Epoch 7, global step 341: 'val_ner_token_f1' was not in top 3
Start to fuse 3 checkpoints via the greedy soup algorithm.
💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
AutoMM has created your model. 🎉🎉🎉
To load the model, use the code below:
```python
from autogluon.multimodal import MultiModalPredictor
predictor = MultiModalPredictor.load("/home/ci/autogluon/docs/tutorials/multimodal/text_prediction/tmp/554fed60d4e945649b83069b89c38b1c-automm_ner")
```
If you are not satisfied with the model, try to increase the training time,
adjust the hyperparameters (https://auto.gluon.ai/stable/tutorials/multimodal/advanced_topics/customization.html),
or post issues on GitHub (https://github.com/autogluon/autogluon/issues).
<autogluon.multimodal.predictor.MultiModalPredictor at 0x7efd181cc5f0>
Evaluation¶
To check the model performance on the test dataset, all you need to do is to call predictor.evaluate(...).
predictor.evaluate(dev_data)
💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
{'hccx': {'precision': np.float64(0.7916967509025271),
'recall': np.float64(0.8596628773030184),
'f1': np.float64(0.8242811501597442),
'number': np.int64(2551)},
'hppx': {'precision': np.float64(0.6859205776173285),
'recall': np.float64(0.6834532374100719),
'f1': np.float64(0.6846846846846847),
'number': np.int64(278)},
'misc': {'precision': np.float64(0.673721340388007),
'recall': np.float64(0.7579365079365079),
'f1': np.float64(0.7133520074696545),
'number': np.int64(504)},
'xh': {'precision': np.float64(0.6531365313653137),
'recall': np.float64(0.7436974789915967),
'f1': np.float64(0.6954813359528488),
'number': np.int64(238)},
'overall_precision': np.float64(0.7572715572715573),
'overall_recall': np.float64(0.8238588630635676),
'overall_f1': np.float64(0.7891630901287553),
'overall_accuracy': 0.8837449267386669}
Prediction and Visualization¶
You can easily obtain the predictions given an input sentence by by calling predictor.predict(...).
output = predictor.predict(dev_data)
visualize_ner(dev_data["text_snippet"].iloc[0], output[0])
💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
家用防尘厨房厨师帽子 HCCX 车间工厂鸭舌 HCCX 工作帽 HCCX 男女食堂餐厅食品 HCCX 卫生帽 HCCX
Now, let’s make predictions on the rabbit toy example.
sentence = "2023年兔年挂件新年装饰品小挂饰乔迁之喜门挂小兔子"
predictions = predictor.predict({'text_snippet': [sentence]})
visualize_ner(sentence, predictions[0])
💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
2023年兔年挂件 HCCX 新年装饰品 HCCX 小挂饰 HCCX 乔迁 MISC 之 HPPX 喜门挂 HCCX 小兔子 HCCX
Other Examples¶
You may go to AutoMM Examples to explore other examples about AutoMM.
Customization¶
To learn how to customize AutoMM, please refer to Customize AutoMM.