search

szxiangjn / world-model-for-language-model

111 stars 5 forks source link

readme

World Model for Language Model

Official code for Language Models Meet World Models: Embodied Experiences Enhance Language Models (NeurIPS 2023). Also check our twitter here.

struct

Setting Up

Install the dependencies by

pip install -r requirements.txt

E2WM Benchmark

struct


All constructed training data from embodied experiences and evaluation data are released in /data/train and /data/eval, respectively.

Data Statistics - Training | Task | Size | |----------------------|------| | Plan Generation | 1659 | | Activity Recognition | 1659 | | Counting | 1000 | | Object Path Tracking | 1000 | - Evaluation | Task | Size | |------------------------------------|------| | Plan Generation | | |   - Vanilla Seen | 175 | |   - Vanilla Unseen | 54 | |   - Confusing Seen | 135 | |   - Confusing UnSeen | 43 | | Houwork QA | 261 | | Negation Housework QA | 162 | | Activity Recognition QA | 549 | | Activity Inference QA | 262 | | Counting QA | 194 | | Object Path Tracking | 200 | | Object Location QA | 200 |

Train & Eval

We compute the fiser matrixs on the sampled 20000 examples from Pile validation set. You can download fisher-matrix-1.3B and fisher-matrix-6B from huggingface model hub, and put them under fisher-matrix directory.

Then go to the scripts directory where you can find all the training and evaluation scripts:

cd scripts

We provide usage examples of GPT-J-6B below. If you want to use GPJ-Neo-1.3B, just replace 6B in the script name with 1.3B.

Train

If you want to train GPT-J-6B, use:

sh run_6B.sh

This script trains GPT-J-6B on a single GPU.

If you want to do distributed training, first run this:

accelerate config --config_file accelerate_config.json

and follow the instructions to set up the config file. (We also provide a sample config file in scripts)

Then, you can simply run:

sh run_6B_multi_gpu.sh

Eval

To do evaluation on QA tasks and generation tasks, run 

sh eval_qa_6B.sh

and 

sh eval_gen_6B.sh

The results will be stored in output/ewc-lora-6B/qa-metric.txt and output/ewc-lora-6B/gen-metric.txt, respectively.

If you want to do distributed evaluation for generation tasks, please modify eval_gen_6B.sh as:

  1. Replace python eval_gen.py with accelerate launch --config_file accelerate_config.json eval_gen.py
  2. Remove export CUDA_VISIBLE_DEVICES=0

This is same as the change from run_6B.sh to run_6B_multi_gpu.sh.

Citation

@article{xiang2023language,
  title={Language Models Meet World Models: Embodied Experiences Enhance Language Models},
  author={Xiang, Jiannan and Tao, Tianhua and Gu, Yi and Shu, Tianmin and Wang, Zirui and Yang, Zichao and Hu, Zhiting},
  journal={arXiv preprint arXiv:2305.10626},
  year={2023}
}