# Mastering Symbolic Operations: Augmenting Language Models with Compiled Neural Networks

### Install 

```
cd ..
git clone https://github.com/deepmind/tracr
cd tracr
pip install .

cd ..
cd code
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)
```

### Use AutoCoNN to create your CoNN

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).

```python
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)
```







### Use CoNN from huggingface
You can load our CoNN saved in Pytorch format through Huggingface. Here is an example of a REVERSE CoNN that can reverse the input text.
```python
from NeuralCom.CoNN.modeling_conn import CoNNModel
from NeuralCom.CoNN import Tokenizer


model = CoNNModel.from_pretrained('model/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']]
```


### One-click implementation for Neural-Comprehension
You can try the differences between our model and the Vanilla model in the coin flip task. We provide a one-click combination method, in which 'WENGSYX/PLM_T5_Base_coin_flip' comes from a model pre-trained on the coin flip task, and CoNN is the chosen Parity CoNN. Through our provided NCModelForCoinFlip, you can combine these two models in one click within the Neural Comprehension framework.

```python
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('model/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. Musk flip the coin. Alex 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}")
```


#### Huggingface Model

In each link, we provide detailed instructions on how to use the CoNN model.

| Model Name  | Model Size |
| ----------- | ---------- |
| Parity      | 2.2M       |
| Reverse     | 4.3M       |
| Last Letter | 62.6K      |
| Copy        | 8.8K       |
| Add_Carry   | 117K       |
| Sub_Carry   | 117K       |

###### If you have also created some amazing CoNN, you are welcome to share them publicly with us.


## 🌱 Neural-Comprehension's Roadmap 🌱


Our future plan includes but not limited to :
- [x] One-click implementation of integration between CoNN and PLM (huggingface)


###### 
