Authors: Yixuan Weng, Minjun Zhu, Fei Xia, Bin Li, Shizhu He, Kang Liu, Jun Zhao 😎
[Contact] If you have any questions, feel free to contact me via (wengsyx@gmail.com).
[Openreview] You can view the peer review comments during the review period or raise an open discussion here.
This repository contains code, models, and other related resources of our paper "Mastering Symbolic Operations: Augmenting Language Models with Compiled Neural Networks".
- [2024/03/08] We release the carame ready verison for this paper!
- [2024/01/16] The Paper accepted in ICLR 2024 Poster!
- [2023/06/21] We release the accompanying project Lingo!
- [2023/05/19] We have supported the one-click implementation of integration between CoNN and PLM!
- [2023/05/18] We have published the "paper-v2"!
- [2023/04/04] We have used huggingface to release the weight of the CoNN model!
- [2023/04/04] We have released the code for AutoCoNN!
- [2023/04/03] We have published the paper!
- [2023/03/26] We created the Github library!
git clone https://github.com/deepmind/tracr
cd tracr
pip3 install .
cd ..
git clone https://github.com/WENGSYX/Neural-Comprehension
cd Neural-Comprehension
pip install .
To run neural comprehension, you need to install PyTorch
, Transformers
, jax
, and tracr
.
# https://beta.openai.com/account/api-keys
export OPENAI_API_KEY=(YOUR OPENAI API KEY)
Please note that setting an OpenAI Key is required to use AutoCoNN (but not necessary if you're just experimenting with neural cognition and CoNN models).
from NeuralCom.AutoCoNN import AutoCoNN
INSTRUCT = 'Create an SOp that is the last letter of a word'
VOCAB = ['a','b','c','d','e','f','g']
EXAMPLE = [[['a','b','c'],['c','c','c']],[['b','d'],['d','d']]]
auto = AutoCoNN()
model,tokenizer = auto(instruct=INSTRUCT,vocab=VOCAB,example=EXAMPLE)
from NeuralCom.CoNN.modeling_conn import CoNNModel
from NeuralCom.CoNN import Tokenizer
model = CoNNModel.from_pretrained('WENGSYX/CoNN_Reverse')
tokenizer = Tokenizer(model.config.input_encoding_map, model.config.output_encoding_map,model.config.max_position_embeddings)
output = model(tokenizer('r e v e r s e').unsqueeze(0))
print(tokenizer.decode(output.argmax(2)))
>>> [['bos', 'e', 's', 'r', 'e', 'v', 'e', 'r']]
from transformers import AutoModel,AutoTokenizer,AutoModelForSeq2SeqLM
from NeuralCom.CoNN.modeling_conn import CoNNModel
from NeuralCom.CoNN import Tokenizer as CoNNTokenizer
from NeuralCom.Model import NCModelForCoinFlip
PLM = AutoModelForSeq2SeqLM.from_pretrained('WENGSYX/PLM_T5_Base_coin_flip')
CoNN = CoNNModel.from_pretrained('WENGSYX/CoNN_Parity')
PLMTokenizer = AutoTokenizer.from_pretrained('WENGSYX/PLM_T5_Base_coin_flip')
CoNNTokenizer = CoNNTokenizer(CoNN.config.input_encoding_map, CoNN.config.output_encoding_map,CoNN.config.max_position_embeddings)
neural_comprehension = NCModelForCoinFlip(PLM, CoNN, PLMTokenizer, CoNNTokenizer).to('cuda:0')
input_text = "A coin is heads up. Aaron flips the coin. Julius does not flip the coin. Yixuan Weng flip the coin. Minjun Zhu does not flip the coin. Is the coin still heads up?"
input_tokens_PLM = PLMTokenizer.encode(input_text, return_tensors='pt')
NC_generated_output = neural_comprehension.generate(input_tokens_PLM.to('cuda:0'))
NC_generated_text = PLMTokenizer.decode(NC_generated_output, skip_special_tokens=True)
Vanilla_generated_output = PLM.generate(input_tokens_PLM.to('cuda:0'))
Vanilla_generated_text = PLMTokenizer.decode(NC_generated_output, skip_special_tokens=True)
print(f"Vanilla Output: {Vanilla_generated_text}")
print(f"Neural Comprehension Output: {NC_generated_text}")
In each link, we provide detailed instructions on how to use the CoNN model.
Model Name | Model Size | Model Address |
---|---|---|
Parity | 2.2M | [link] |
Reverse | 4.3M | [link] |
Last Letter | 62.6K | [link] |
Copy | 8.8K | [link] |
Add_Carry | 117K | [link] |
Sub_Carry | 117K | [link] |
Our future plan includes but not limited to :
- One-click implementation of integration between CoNN and PLM (huggingface)
- Demo Presentation
@inproceedings{
weng2024mastering,
title={Mastering Symbolic Operations: Augmenting Language Models with Compiled Neural Networks},
author={Yixuan Weng and Minjun Zhu and Fei Xia and Bin Li and Shizhu He and Kang Liu and Jun Zhao},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=9nsNyN0vox}
}