diff --git a/examples/tutorial/2_with_bmtrain.py b/examples/tutorial/2_with_bmtrain.py index 45b0d99..bfb677a 100644 --- a/examples/tutorial/2_with_bmtrain.py +++ b/examples/tutorial/2_with_bmtrain.py @@ -1,5 +1,8 @@ # adapted from https://github.com/OpenBMB/ModelCenter/blob/main/examples/bert/finetune_bert.py +# For ModelCenter, `pip install model_center >= 1.0.0` +# For BMTrain, `git clone https://github.com/OpenBMB/BMTrain.git` and `python3 setup.py install` locally, as it has not been released currently. + import time import os @@ -56,7 +59,8 @@ def get_model(args): "WiC" : 2, } model = BertModel(args, num_types[args.dataset_name]) - # od.Visualization(model).structure_graph() + from bigmodelvis import Visualization + Visualization(model).structure_graph() if args.delta_type == "lora": delta_model = LoraModel(backbone_model=model, modified_modules=['project_q', 'project_k'], backend='bmt')