fix: config

.gitignore

@@ -1,8 +1,3 @@
-*.arrow
-*.swp
-# ignore knn index
-gpt4all/index/**/
-.DS_Store
 *.pkl
 ckpts*
 .deepspeed_env
@@ -169,9 +164,4 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-
-
-# vs code
-.vscode
-*.bin
+#.idea/

.gitmodules

@@ -1,3 +1,6 @@
+[submodule "transformers"]
+	path = transformers
+	url = https://github.com/huggingface/transformers.git
 [submodule "peft"]
 	path = peft
 	url = https://github.com/huggingface/peft.git

GPT-J_MAP.md

@@ -1,17 +0,0 @@
-# Inference on Training Data
-
-
-## Run Inference
-
-```bash
-torchrun --master_port=29085 --nproc-per-node 8 inference.py --config=configs/inference/gptj.yaml
-```
-
-
-## Visualizations
-
-```bash
-python build_map.py
-```
- 
-will build a map in `Atlas`, one using the internal clustering algorithm provided by Nomic and one using the embeddings generated by the finetuned model.

LICENSE.txt

@@ -1,19 +0,0 @@
-Copyright (c) 2023 Nomic, Inc.
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

README.md

@@ -1,118 +1,13 @@
 <h1 align="center">GPT4All</h1>
-<p align="center">Demo, data, and code to train open-source assistant-style large language model based on GPT-J and LLaMa</p>
-<p align="center">
-<a href="https://static.nomic.ai/gpt4all/2023_GPT4All-J_Technical_Report_2.pdf">:green_book: Technical Report 2: GPT4All-J </a>
-</p>
+<p align="center">Demo, data and code to train an assistant-style large language model with ~800k GPT-3.5-Turbo Generations based on LLaMa</p>
 
 <p align="center">
-<a href="https://s3.amazonaws.com/static.nomic.ai/gpt4all/2023_GPT4All_Technical_Report.pdf">:green_book: Technical Report 1: GPT4All</a>
+<a href="https://s3.amazonaws.com/static.nomic.ai/gpt4all/2023_GPT4All_Technical_Report.pdf">:green_book: Technical Report</a>
 </p>
-
 <p align="center">
-<a href="https://github.com/nomic-ai/pyllamacpp">:snake: Official Python Bindings</a>
+<a href="https://discord.gg/kvmy6dQB">Discord</a>
 </p>
 
-<p align="center">
-<a href="https://github.com/nomic-ai/gpt4all-ts">:computer: Official Typescript Bindings</a>
-</p>
-
-<p align="center">
-<a href="https://github.com/nomic-ai/gpt4all-ui">:speech_balloon: Official Web Chat Interface</a>
-</p>
-
-<p align="center">
-<a href="https://github.com/nomic-ai/gpt4all-chat">:speech_balloon: Official Chat Interface</a>
-</p>
-
-<p align="center">
-<a href="https://python.langchain.com/en/latest/modules/models/llms/integrations/gpt4all.html">🦜️🔗 Official Langchain Backend</a> 
-</p>
-
-
-<p align="center">
-<a href="https://discord.gg/mGZE39AS3e">Discord</a>
-</p>
-
-
-
-
-<p align="center">
-GPT4All is made possible by our compute partner <a href="https://www.paperspace.com/">Paperspace</a>.
-</p>
-
-
-
-## GPT4All-J: An Apache-2 Licensed GPT4All Model
-![gpt4all-j-demo](https://user-images.githubusercontent.com/13879686/231876409-e3de1934-93bb-4b4b-9013-b491a969ebbc.gif)
-
-Run on an M1 Mac (not sped up!)
-
-
-### GPT4All-J Chat UI Installers
-Installs a native chat-client with auto-update functionality that runs on your desktop with the GPT4All-J model baked into it.
-
-[Mac/OSX](https://gpt4all.io/installers/gpt4all-installer-darwin.dmg)
-
-[Windows](https://gpt4all.io/installers/gpt4all-installer-win64.exe)
-
-[Ubuntu](https://gpt4all.io/installers/gpt4all-installer-linux.run)
-
-If you have older hardware that only supports avx and not avx2 you can use these.
-
-[Mac/OSX - avx-only](https://gpt4all.io/installers/gpt4all-installer-darwin-avx-only.dmg)
-
-[Windows - avx-only](https://gpt4all.io/installers/gpt4all-installer-win64-avx-only.exe)
-
-[Ubuntu - avx-only](https://gpt4all.io/installers/gpt4all-installer-linux-avx-only.run)
-
-These files are not yet cert signed by Windows/Apple so you will see security warnings on initial installation. We did not want to delay release while waiting for their process to complete.
-
-Find the most up-to-date information on the [GPT4All Website](https://gpt4all.io/)
-
-### Raw Model
-[ggml Model Download Link](https://gpt4all.io/models/ggml-gpt4all-j.bin)
-
-Note this model is only compatible with the C++ bindings found [here](https://github.com/nomic-ai/gpt4all-chat). It will not work with any existing llama.cpp bindings as we had to do a large fork of llama.cpp. GPT4All will support the ecosystem around this new C++ backend going forward.
-
-Python bindings are imminent and will be integrated into this [repository](https://github.com/nomic-ai/pyllamacpp). Stay tuned on the [GPT4All discord](https://discord.gg/mGZE39AS3e) for updates.
-
-## Training GPT4All-J
-
-Please see [GPT4All-J Technical Report](https://static.nomic.ai/gpt4all/2023_GPT4All-J_Technical_Report_2.pdf) for details.
-
-### GPT4All-J Training Data
-
-- We are releasing the curated training data for anyone to replicate GPT4All-J here: [GPT4All-J Training Data](https://huggingface.co/datasets/nomic-ai/gpt4all-j-prompt-generations)
-   - [Atlas Map of Prompts](https://atlas.nomic.ai/map/gpt4all-j-prompts-curated)
-   - [Atlas Map of Responses](https://atlas.nomic.ai/map/gpt4all-j-response-curated)
-   
-We have released updated versions of our `GPT4All-J` model and training data. 
-
-- `v1.0`: The original model trained on the v1.0 dataset
-- `v1.1-breezy`: Trained on a filtered dataset where we removed all instances of AI language model
-- `v1.2-jazzy`: Trained on a filtered dataset where we also removed instances like I'm sorry, I can't answer... and AI language model
-
-The [models](https://huggingface.co/nomic-ai/gpt4all-j) and [data](https://huggingface.co/datasets/nomic-ai/gpt4all-j-prompt-generations) versions can be specified by passing a `revision` argument.
-
-For example, to load the `v1.2-jazzy` model and dataset, run:
-
-```python
-from datasets import load_dataset
-from transformers import AutoModelForCausalLM
-
-dataset = load_dataset("nomic-ai/gpt4all-j-prompt-generations", revision="v1.2-jazzy")
-model = AutoModelForCausalLM.from_pretrained("nomic-ai/gpt4all-j-prompt-generations", revision="v1.2-jazzy")
-```
-
-### GPT4All-J Training Instructions
-
-```bash
-accelerate launch --dynamo_backend=inductor --num_processes=8 --num_machines=1 --machine_rank=0 --deepspeed_multinode_launcher standard --mixed_precision=bf16  --use_deepspeed --deepspeed_config_file=configs/deepspeed/ds_config_gptj.json train.py --config configs/train/finetune_gptj.yaml
-```
-
-
-# Original GPT4All Model (based on GPL Licensed LLaMa)
-
 
 
 ![gpt4all-lora-demo](https://user-images.githubusercontent.com/13879686/228352356-de66ca7a-df70-474e-b929-2e3656165051.gif)
@@ -121,115 +16,31 @@ Run on M1 Mac (not sped up!)
 
 # Try it yourself
 
-Here's how to get started with the CPU quantized GPT4All model checkpoint:
+Download the CPU quantized gpt4all model checkpoint: [gpt4all-lora-quantized.bin](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-quantized.bin).
 
-1. Download the `gpt4all-lora-quantized.bin` file from [Direct Link](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-quantized.bin) or [[Torrent-Magnet]](https://tinyurl.com/gpt4all-lora-quantized).
-2. Clone this repository, navigate to `chat`, and place the downloaded file there.
-3. Run the appropriate command for your OS:
-   - M1 Mac/OSX: `cd chat;./gpt4all-lora-quantized-OSX-m1`
-   - Linux: `cd chat;./gpt4all-lora-quantized-linux-x86`
-   - Windows (PowerShell): `cd chat;./gpt4all-lora-quantized-win64.exe`
-   - Intel Mac/OSX: `cd chat;./gpt4all-lora-quantized-OSX-intel`
 
-For custom hardware compilation, see our [llama.cpp](https://github.com/zanussbaum/gpt4all.cpp) fork.
+Clone this repository down and place the quantized model in the `chat` directory and start chatting by running:
 
------------
-Find all compatible models in the GPT4All Ecosystem section.
+- `cd chat;./gpt4all-lora-quantized-OSX-m1` on M1 Mac/OSX
+- `cd chat;./gpt4all-lora-quantized-linux-x86` on Linux
+- `cd chat;./gpt4all-lora-quantized-win64.exe` on Windows (PowerShell)
+- `cd chat;./gpt4all-lora-quantized-OSX-intel` on Intel Mac/OSX
 
-[Secret Unfiltered Checkpoint](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-unfiltered-quantized.bin) - [[Torrent]](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-unfiltered-quantized.bin.torrent)
+To compile for custom hardware, see our fork of the [Alpaca C++](https://github.com/zanussbaum/gpt4all.cpp) repo.
 
-This model had all refusal to answer responses removed from training. Try it with:
-- M1 Mac/OSX: `cd chat;./gpt4all-lora-quantized-OSX-m1 -m gpt4all-lora-unfiltered-quantized.bin`
-- Linux: `cd chat;./gpt4all-lora-quantized-linux-x86 -m gpt4all-lora-unfiltered-quantized.bin`
-- Windows (PowerShell): `cd chat;./gpt4all-lora-quantized-win64.exe -m gpt4all-lora-unfiltered-quantized.bin`
-- Intel Mac/OSX: `cd chat;./gpt4all-lora-quantized-OSX-intel -m gpt4all-lora-unfiltered-quantized.bin`
------------
 Note: the full model on GPU (16GB of RAM required) performs much better in our qualitative evaluations.
 
-# Python Client
-## CPU Interface
-To run GPT4All in python, see the new [official Python bindings](https://github.com/nomic-ai/pyllamacpp).
-
-The old bindings are still available but now deprecated. They will not work in a notebook environment.
-To get running using the python client with the CPU interface, first install the [nomic client](https://github.com/nomic-ai/nomic) using `pip install nomic`
-Then, you can use the following script to interact with GPT4All:
-```
-from nomic.gpt4all import GPT4All
-m = GPT4All()
-m.open()
-m.prompt('write me a story about a lonely computer')
-```
-
-## GPU Interface
-There are two ways to get up and running with this model on GPU.
-The setup here is slightly more involved than the CPU model.
-1. clone the nomic client [repo](https://github.com/nomic-ai/nomic) and run `pip install .[GPT4All]` in the home dir.
-2. run `pip install nomic` and install the additional deps from the wheels built [here](https://github.com/nomic-ai/nomic/tree/main/bin)
-
-Once this is done, you can run the model on GPU with a script like the following:
-```
-from nomic.gpt4all import GPT4AllGPU
-m = GPT4AllGPU(LLAMA_PATH)
-config = {'num_beams': 2,
-          'min_new_tokens': 10,
-          'max_length': 100,
-          'repetition_penalty': 2.0}
-out = m.generate('write me a story about a lonely computer', config)
-print(out)
-```
-Where LLAMA_PATH is the path to a Huggingface Automodel compliant LLAMA model.
-Nomic is unable to distribute this file at this time.
-We are working on a GPT4All that does not have this limitation right now.
-
-You can pass any of the [huggingface generation config params](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationConfig) in the config.
-
-# GPT4All Compatibility Ecosystem
-Edge models in the GPT4All Ecosystem. Please PR as the [community grows](https://huggingface.co/models?sort=modified&search=4bit).
-Feel free to convert this to a more structured table.
-
-- [gpt4all](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-quantized.bin) [[MD5 Signature](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-quantized.bin.md5)]
-   - [gpt4all-ggml-converted](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-quantized-ggml.bin) [[MD5 Signature](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-quantized-ggml.bin.md5)]
-- [gpt4all-unfiltered](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-unfiltered-quantized.bin) [[MD5 Signature](https://the-eye.eu/public/AI/models/nomic-ai/gpt4all/gpt4all-lora-unfiltered-quantized.bin.md5)]
-- [ggml-vicuna-7b-4bit](https://huggingface.co/eachadea/ggml-vicuna-7b-4bit)
-- [vicuna-13b-GPTQ-4bit-128g](https://huggingface.co/anon8231489123/vicuna-13b-GPTQ-4bit-128g)
-- [LLaMa-Storytelling-4Bit](https://huggingface.co/GamerUntouch/LLaMa-Storytelling-4Bit)
-- [Alpaca Native 4bit](https://huggingface.co/Sosaka/Alpaca-native-4bit-ggml/tree/main)
-
-
-# Roadmap
-## Short Term
- - <span style="color:green">(Done)</span> Train a GPT4All model based on GPTJ to alleviate llama distribution issues.
- - <span style="color:green">(Done)</span> Create improved CPU and GPU interfaces for this model.
- - <span style="color:green">(Done)</span> [Integrate llama.cpp bindings](https://github.com/nomic-ai/pyllamacpp)
- - <span style="color:green">(Done)</span> [Create a good conversational chat interface for the model.](https://github.com/nomic-ai/gpt4all-ui)
- - <span style="color:green">(Done)</span> [Allow users to opt in and submit their chats for subsequent training runs](https://github.com/nomic-ai/gpt4all-ui)
-
-## Medium Term
- - <span style="color:red">(NOT STARTED)</span> Integrate GPT4All with [Atlas](https://atlas.nomic.ai) to allow for document retrieval.
-   - BLOCKED by GPT4All based on GPTJ
- - <span style="color:red">(Done)</span> Integrate GPT4All with Langchain.
- - <span style="color:green">(IN PROGRESS)</span> Build easy custom training scripts to allow users to fine tune models.
-
-## Long Term
- - <span style="color:red">(NOT STARTED)</span> Allow anyone to curate training data for subsequent GPT4All releases using Atlas.
- - <span style="color:green">(IN PROGRESS)</span> Democratize AI. 
-
 # Reproducibility
 
-Trained Model Weights:
+Trained LoRa Weights:
 - gpt4all-lora (four full epochs of training):  https://huggingface.co/nomic-ai/gpt4all-lora
 - gpt4all-lora-epoch-2 (three full epochs of training) https://huggingface.co/nomic-ai/gpt4all-lora-epoch-2
-- gpt4all-j (one full epoch of training) (https://huggingface.co/nomic-ai/gpt4all-j)
-- gpt4all-j-lora (one full epoch of training) (https://huggingface.co/nomic-ai/gpt4all-j-lora)
 
 Raw Data:
-- [Training Data Without P3](https://huggingface.co/datasets/nomic-ai/gpt4all_prompt_generations)
+- [Training Data Without P3](https://s3.amazonaws.com/static.nomic.ai/gpt4all/2022_03_27/gpt4all_curated_data_without_p3_2022_03_27.tar.gz)
   - Explorer: https://atlas.nomic.ai/map/gpt4all_data_clean_without_p3
-- [Full Dataset with P3](https://huggingface.co/datasets/nomic-ai/gpt4all_prompt_generations_with_p3)
+- [Full Dataset with P3](https://s3.amazonaws.com/static.nomic.ai/gpt4all/2022_03_27/gpt4all_curated_data_full_2022_03_27.tar.gz)
   - Explorer: https://atlas.nomic.ai/map/gpt4all_data_clean
-- [GPT4All-J Dataset](https://huggingface.co/datasets/nomic-ai/gpt4all-j-prompt-generations)
-   -  Explorer Indexed on Prompts: https://atlas.nomic.ai/map/gpt4all-j-prompts-curated
-   -  Exporer Indexed on Responses: https://atlas.nomic.ai/map/gpt4all-j-response-curated
 
 We are not distributing a LLaMa 7B checkpoint.
 
@@ -239,16 +50,18 @@ You can reproduce our trained model by doing the following:
 
 Clone the repo
 
-```
-git clone --recurse-submodules https://github.com/nomic-ai/gpt4all.git
-git submodule update --init
-```
+`git clone --recurse-submodules https://github.com/nomic-ai/gpt4all.git`
+
+`git submodule configure && git submodule update`
 
 Setup the environment
 
 ```
 python -m pip install -r requirements.txt
 
+cd transformers
+pip install -e . 
+
 cd ../peft
 pip install -e .
 ```
@@ -265,10 +78,6 @@ accelerate launch --dynamo_backend=inductor --num_processes=8 --num_machines=1 -
 python generate.py --config configs/generate/generate.yaml --prompt "Write a script to reverse a string in Python"
 ```
 
-## Need Help?
-
-Join the <a href="https://discord.gg/E7q6fp9bmx"> Discord </a> and ask for help in `#gpt4all-help`
-
 # Sample Generations
 
 ### Provide instructions for the given exercise. Leg Raises
@@ -295,8 +104,7 @@ Join the <a href="https://discord.gg/E7q6fp9bmx"> Discord </a> and ask for help
 
 ### Come up with an interesting idea for a new movie plot. Your plot should be described with a title and a summary.
 
-> 1. The title of your movie plot should be "The Last Stand". 
-> 2. The summary of your movie plot should be: A group of mercenaries are hired to protect a small town from a dangerous criminal who is on the run. However, the town's sheriff is not willing to let them take over and decides to take matters into his own hands.
+>1. The title of your movie plot should be "The Last Stand". 2. The summary of your movie plot should be: A group of mercenaries are hired to protect a small town from a dangerous criminal who is on the run. However, the town's sheriff is not willing to let them take over and decides to take matters into his own hands.
 
 ### Reverse a string in python.
 > In Python, you can reverse a list or tuple by using the reversed() function on it. Here's an example of how to use this method with strings: 
@@ -306,43 +114,42 @@ Join the <a href="https://discord.gg/E7q6fp9bmx"> Discord </a> and ask for help
 >   print(reversed_str)     # Output will be reversed version of the inputted string.
 
 ### List 10 dogs.
-> 1. Labrador Retriever
-> 2. Golden Retriever
-> 3. Beagle
-> 4. Poodle
-> 5. German Shepherd
-> 6. Bulldog
-> 7. Rottweiler
-> 8. Boxer
-> 9. Shih Tzu
-> 10. Pug
+>1. Labrador Retriever
+>2. Golden Retriever
+>3. Beagle
+>4. Poodle
+>5. German Shepherd
+>6. Bulldog
+>7. Rottweiler
+>8. Boxer
+>9. Shih Tzu
+>10. Pug
 
-### Write me a poem about the fall of Julius Caesar into a caesar salad in iambic pentameter. 
-
-> The fall of Julius Caesar into a caesar salad<br>
-> In iambic pentameter,<br>
-> The mighty Roman emperor<br>
-> Fell into a salad of lettuce and croutons,<br>
-> His empire crumbling around him,<br>
-> As he was devoured by the hungry diners.<br>
-> The once mighty emperor<br>
-> Was now just a salad topping,<br>
-> His legacy reduced to a mere garnish,<br>
-> As he was consumed by the masses.<br>
-> The fall of Julius Caesar into a caesar salad,<br>
-> In iambic pentameter,<br>
-> Was a tragic end to a great man,<br>
-> Who had once ruled the world with an iron fist.<br>
-> But now, he was just a salad topping,<br>
-> His legacy reduced to a mere garnish,<br>
-> As he was consumed by the masses.
+### Write me a poem about the fall of Julius Ceasar into a ceasar salad in iambic pentameter. 
+>The fall of Julius Ceasar into a ceasar salad
+>In iambic pentameter,
+>The mighty Roman emperor
+>Fell into a salad of lettuce and croutons,
+>His empire crumbling around him,
+>As he was devoured by the hungry diners.
+>The once mighty emperor
+>Was now just a salad topping,
+>His legacy reduced to a mere garnish,
+>As he was consumed by the masses.
+>The fall of Julius Ceasar into a ceasar salad,
+>In iambic pentameter,
+>Was a tragic end to a great man,
+>Who had once ruled the world with an iron fist.
+>But now, he was just a salad topping,
+>His legacy reduced to a mere garnish,
+>As he was consumed by the masses.
     
 ### What is a three word topic describing the following keywords: baseball, football, soccer: 
-> Sports, athletics, games
+>Sports, athletics, games
     
-## Citation
 
-If you utilize this repository, models or data in a downstream project, please consider citing it with:
+
+If you utilize this reposistory, models or data in a downstream project, please consider citing it with:
 ```
 @misc{gpt4all,
   author = {Yuvanesh Anand and Zach Nussbaum and Brandon Duderstadt and Benjamin Schmidt and Andriy Mulyar},
@@ -353,3 +160,7 @@ If you utilize this repository, models or data in a downstream project, please c
   howpublished = {\url{https://github.com/nomic-ai/gpt4all}},
 }
 ```
+
+### Alternative Download Locations
+#### gpt4all-lora-quantized.bin Backup Torrent Link
+magnet:?xt=urn:btih:1F11A9691EE06C18F0040E359361DCA0479BCB5A&dn=gpt4all-lora-quantized.bin&tr=udp%3A%2F%2Ftracker.opentrackr.org%3A1337%2Fannounce&tr=udp%3A%2F%2Fopentracker.i2p.rocks%3A6969%2Fannounce

TRAINING_LOG.md

@@ -23,7 +23,7 @@ We used the initial parameters:
 | Weight decay   | 0     |
 | Warmup Steps   | 100   |
 
-We randomly shuffle and set aside 5% of the data for validation.
+We randomly shuffle and set aside %5 of the data for validation.
 
 We had an initial bug in logging the training loss but we noticed a decrease in validation loss.
 
@@ -160,7 +160,7 @@ We realized that we had two bugs however:
 - We accidentally duplicated data and effectively trained for 2 epochs instead of 1
 - We added an eos token to every sequence, even those that we truncated (e.g. long code that exceeds the 1024).
 
-## Conditional EOS and 1 Epoch
+## Conditonal EOS and 1 Epoch
 
 Using the same parameters, we then trained a model using a "conditional" eos token where we only add an `eos` when the inputs are less than the maximum sequence length for one epoch.
 
@@ -235,49 +235,3 @@ Taking inspiration from [the Alpaca Repo](https://github.com/tatsu-lab/stanford_
 Comparing our model LoRa to the [Alpaca LoRa](https://huggingface.co/tloen/alpaca-lora-7b), our model has lower perplexity. Qualitatively, training on 3 epochs performed the best on perplexity as well as qualitative examples. 
 
 We tried training a full model using the parameters above, but found that during the second epoch the model diverged and samples generated post training were worse than the first epoch. 
-
-
-## GPT-J Training
-
-### Model Training Divergence
-
-We trained multiple [GPT-J models](https://huggingface.co/EleutherAI/gpt-j-6b) with varying success. We found that training the full model lead to diverged post epoch 1. ![](figs/overfit-gpt-j.png)
-
-
-We release the checkpoint after epoch 1.
-
-
-Using Atlas, we extracted the embeddings of each point in the dataset and calculated the loss per sequence. We then uploaded [this to Atlas](https://atlas.nomic.ai/map/gpt4all-j-post-epoch-1-embeddings) and noticed that the higher loss items seem to cluster. On further inspection, the highest density clusters seemded to be of prompt/response pairs that asked for creative-like generations such as `Generate a story about ...` ![](figs/clustering_overfit.png)
-
-
-
-### GPT4All-J Hyperparameters
-
-We varied learning rate, learning rate schedule, and weight decay following suggestions from the [original GPT-J codebase](https://github.com/kingoflolz/mesh-transformer-jax/blob/master/howto_finetune.md) but found no real performance difference (qualitatively or quantitatively) when varying these parameters.
-
-
-
-The final model was trained using the following hyperparameters with a linear warmup followed by constant learning rate:
-
-| Hyperparameter | Value |
-|----------------|-------|
-| Per Device BS  | 32    |
-| Global BS      | 256   |
-| Learning rate  | 2e-5  |
-| Epochs         | 2     |
-| Max length     | 1024  |
-| Weight decay   | 0     |
-| Warmup Steps   | 500   |
-
-
-The LoRA model was trained using using the following hyperparameters with a linear warmup followed by constant learning rate: 
-
-| Hyperparameter | Value |
-|----------------|-------|
-| Per Device BS  | 4     |
-| Global BS      | 32    |
-| Learning rate  | 2e-5  |
-| Epochs         | 2     |
-| Max length     | 1024  |
-| Weight decay   | 0     |
-| Warmup Steps   | 500   |

build_map.py

@@ -1,54 +0,0 @@
-import numpy as np
-from nomic import atlas
-import glob
-from tqdm import tqdm
-from datasets import load_dataset, concatenate_datasets
-from sklearn.decomposition import PCA
-
-files = glob.glob("inference/*.jsonl")
-print(files)
-df = concatenate_datasets([load_dataset("json", data_files=file, split="train") for file in tqdm(files)])
-
-print(len(df))
-print(df)
-
-df = df.map(lambda example: {"inputs": [prompt + "\n" + response for prompt, response in zip(example["prompt"], example["response"])]},
-            batched=True,
-            num_proc=64)
-
-df = df.map(lambda example: {"trained_on": [int(t) for t in example["is_train"]]},
-                batched=True,
-                num_proc=64)
-                
-df = df.remove_columns("is_train")
-
-text = df.remove_columns(["labels", "input_ids", "embeddings"])
-
-text_df = [text[i] for i in range(len(text))]
-
-atlas.map_text(text_df, indexed_field="inputs",
-               name="CHANGE ME!",
-               colorable_fields=["source", "loss", "trained_on"],
-               reset_project_if_exists=True,
-               )
-
-# index is local to train/test split, regenerate
-data = df.remove_columns(["labels", "input_ids", "index"])
-data = data.add_column("index", list(range(len(data))))
-# max embed dim is 2048 for now
-# note! this is slow in pyarrow/hf datasets
-embeddings = np.array(data["embeddings"])
-print("embeddings shape:", embeddings.shape)
-embeddings = PCA(n_components=2048).fit_transform(embeddings)
-
-data = data.remove_columns(["embeddings"])
-columns = data.to_pandas().to_dict("records") 
-
-atlas.map_embeddings(embeddings,
-                     data=columns,
-                     id_field="index",
-                     name="CHANGE ME!",
-                     colorable_fields=["source", "loss", "trained_on"],
-                     build_topic_model=True,
-                     topic_label_field="inputs",
-                     reset_project_if_exists=True,)

configs/deepspeed/ds_config_gptj.json

@@ -24,16 +24,5 @@
 	  "allgather_partitions": true,
 	  "allgather_bucket_size": 5e8,
 	  "contiguous_gradients": true
-	},
-	"optimizer": {
-		"type": "AdamW",
-		"params": {
-		  "lr": "auto",
-		  "betas": [
-			0.9,
-			0.999
-		  ],
-		  "eps": 1e-08
-		}
-	  }
-}
+	}
+  }

configs/deepspeed/ds_config_pythiaseek.json

@@ -1,39 +0,0 @@
-{
-	"train_batch_size": "auto",
-	"gradient_accumulation_steps": "auto",
-	"train_micro_batch_size_per_gpu": "auto",
-	"fp16": {
-	  "enabled": "auto",
-	  "min_loss_scale": 1,
-	  "loss_scale_window": 1000,
-	  "hysteresis": 2,
-	  "initial_scale_power": 32
-	},
-	"bf16": {
-		"enabled": "auto"
-	},
-	"gradient_clipping": 1.0,
-	"zero_optimization": {
-	  "stage": 2,
-	  "offload_param": {
-		"device": "none"
-	  },
-	  "offload_optimizer": {
-		"device": "none"
-	  },
-	  "allgather_partitions": true,
-	  "allgather_bucket_size": 5e8,
-	  "contiguous_gradients": true
-	},
-	"optimizer": {
-		"type": "AdamW",
-		"params": {
-		  "lr": "auto",
-		  "betas": [
-			0.9,
-			0.999
-		  ],
-		  "eps": 1e-08
-		}
-	  }
-}

configs/eval/evaluate_gpt4all_jr.yaml

@@ -1,18 +0,0 @@
-# model/tokenizer
-model_name: "/home/paperspace/gpt4all/gpt4all/train/ckpts/epoch_2"
-tokenizer_name: "EleutherAI/gpt-j-6B"
-version: null
-gradient_checkpointing: true
-save_name: "nomic-ai/gpt-jr"
-encoder_dim: 384
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "/home/paperspace/gpt4all/gpt4all/index/squad_supplemented_validation" 
-max_length: 1024
-batch_size: 32
-pct_test: 0.05
-q_column: "question"
-a_column: "answers"
-encoder_column: "neighbor_embeddings"

configs/eval/evaluate_lethe.yaml

@@ -1,23 +0,0 @@
-# model/tokenizer
-model_name: "/home/paperspace/gpt4all/gpt4all/train/ckpts/mem_attn/step_1000"
-tokenizer_name: "EleutherAI/pythia-1b"
-version: null
-gradient_checkpointing: false
-memory_attn_layer: 12
-
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "/home/paperspace/gpt4all/gpt4all/inference/synth_data_combined_174"
-# dataset_path: "/home/paperspace/gpt4all/gpt4all/index/squad_supplemented_validation" 
-max_length: 1024
-batch_size: 1
-pct_test: 0.05
-q_column: "question"
-a_column: "answer"
-context_column: "text"
-num_memories_per_index: 2000000
-num_neighbors_to_retrieve: 2
-num_neighbors_to_store: 1 
-mem_chunk_size: 64

configs/eval/evaluate_pythiaseek.yaml

@@ -1,18 +0,0 @@
-# model/tokenizer
-model_name: "/home/paperspace/gpt4all/gpt4all/train/ckpts/learnable_alpha/epoch_0"
-tokenizer_name: "EleutherAI/pythia-1b"
-version: null
-gradient_checkpointing: true
-save_name: "nomic-ai/gpt-jr"
-encoder_dim: 384
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "/home/paperspace/gpt4all/gpt4all/index/squad_supplemented_validation" 
-max_length: 1024
-batch_size: 32
-pct_test: 0.05
-q_column: "question"
-a_column: "answers"
-encoder_column: "neighbor_embeddings"

configs/eval/generate.yaml

@@ -0,0 +1,15 @@
+# model/tokenizer
+model_name: # update with llama 7b 
+tokenizer_name: # update with llama 7b
+lora: true
+lora_path: "nomic-ai/gpt4all-lora"
+
+max_new_tokens: 512
+temperature: 0.001
+prompt: | 
+  #this code prints a string reversed
+  my_string = "hello how are you"
+  print(len(my_string))
+
+
+  My code above does not work. Can you help me?

configs/eval/generate_baseline.yaml

@@ -1,5 +1,17 @@
 # model/tokenizer
-model_name: "zpn/llama-7b"
-tokenizer_name: "zpn/llama-7b"
+model_name: # update with llama model name
+tokenizer_name: # update with llama model name
 lora: true
-lora_path: "tloen/alpaca-lora-7b"
+lora_path: "tloen/alpaca-lora-7b"
+
+
+
+max_new_tokens: 512
+temperature: 0.001
+prompt: | 
+  #this code prints a string reversed
+  my_string = "hello how are you"
+  print(len(my_string))
+
+
+  My code above does not work. Can you help me?

configs/eval/generate_full.yaml

@@ -1,8 +1,7 @@
 # model/tokenizer
-model_name: "nomic-ai/gpt4all-warmup-lr-epoch_1"
-tokenizer_name: "EleutherAI/gpt-j-6b"
-lora: false
-
+model_name: # update
+tokenizer_name: # update
+lora_path: "no-lora"
 
 max_new_tokens: 512
 temperature: 0.001

configs/eval/generate_gpt4all_gptj.yaml

@@ -1,4 +0,0 @@
-# model/tokenizer
-model_name: "nomic-ai/gpt4all-warmup-lr-epoch_0"
-tokenizer_name: "EleutherAI/gpt-j-6b"
-lora: false

configs/eval/generate_gpt4all_gptj_lora.yaml

@@ -1,5 +0,0 @@
-# model/tokenizer
-model_name: "EleutherAI/gpt-j-6b"
-tokenizer_name: "EleutherAI/gpt-j-6B"
-lora: true
-lora_path: "nomic-ai/gpt4all-gptj-lora-epoch_1"

configs/eval/generate_gpt4all_llama_lora.yaml

@@ -1,5 +0,0 @@
-# model/tokenizer
-model_name: "zpn/llama-7b"
-tokenizer_name: "zpn/llama-7b"
-lora: true
-lora_path: "nomic-ai/gpt4all-lora"

configs/eval/generate_large_2.yaml

@@ -0,0 +1,15 @@
+# model/tokenizer
+model_name: # update
+tokenizer_name: # update
+lora: true
+lora_path: # update
+
+max_new_tokens: 512
+temperature: 0.001
+prompt: | 
+  #this code prints a string reversed
+  my_string = "hello how are you"
+  print(len(my_string))
+
+
+  My code above does not work. Can you help me?

configs/eval/generate_large_3.yaml

@@ -0,0 +1,15 @@
+# model/tokenizer
+model_name: # update
+tokenizer_name: # update
+lora: true
+lora_path: # update
+
+max_new_tokens: 512
+temperature: 0.001
+prompt: | 
+  #this code prints a string reversed
+  my_string = "hello how are you"
+  print(len(my_string))
+
+
+  My code above does not work. Can you help me?

configs/generate/generate.yaml

@@ -1,6 +1,6 @@
 # model/tokenizer
-model_name: "zpn/llama-7b"
-tokenizer_name: "zpn/llama-7b"
+model_name: # REPLACE HERE with the base llama model 
+tokenizer_name: # REPLACE HERE with the llama tokenizer
 lora: true
 lora_path: "nomic-ai/gpt4all-lora"
 

configs/generate/generate_gptj_lora.yaml

@@ -1,15 +0,0 @@
-# model/tokenizer
-model_name: "EleutherAI/gpt-j-6b"
-tokenizer_name: "EleutherAI/gpt-j-6b"
-lora: true
-lora_path: "nomic-ai/gpt4all-gptj-lora-epoch_0"
-
-max_new_tokens: 512
-temperature: 0
-prompt: | 
-  #this code prints a string reversed
-  my_string = "hello how are you"
-  print(len(my_string))
-
-
-  My code above does not work. Can you help me?

configs/index/finetune_gptjr.yaml

@@ -1,13 +0,0 @@
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "squad"
-max_length: 1024
-batch_size: 32
-
-#index
-index_path: "/home/paperspace/gpt4all/gpt4all/index/wiki-cohere-index.bin"
-index_database: "nomic-ai/cohere-wiki-sbert"
-index_space: "cosine"
-index_dim: 384
-query_embedding_field: 'question'

configs/inference/gptj.yaml

@@ -1,11 +1,11 @@
 # model/tokenizer
-model_name: "nomic-ai/gpt4all-warmup-lr-epoch_1"
+model_name: "nomic-ai/gpt4all-gptj-multinode-deepspeed-finetuned-epoch_0"
 tokenizer_name: "EleutherAI/gpt-j-6B"
 
 # dataset
 streaming: false
 num_proc: 64
-dataset_path: "nomic-ai/turbo-500k-multi" 
+dataset_path: "data_multiplus" 
 max_length: 1024
 batch_size: 32 
 

configs/inference/synth_data.yaml

@@ -1,19 +0,0 @@
-model_name: "nomic-ai/gpt4all-j"
-revision: "v1.3-groovy"
-tokenizer_name: "nomic-ai/gpt4all-j"
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "nomic-ai/cohere-wiki-sbert"
-batch_size: 32 
-output_path: "synth_qa_pairs"
-
-# generation
-max_new_tokens: 75
-max_generations: 200000
-
-save_every: 1000
-
-
-seed: 42

configs/train/finetune.yaml

@@ -2,7 +2,7 @@
 model_name: # add model here
 tokenizer_name: # add model here
 gradient_checkpointing: true
-save_name: # CHANGE 
+save_name: "nomic-ai/gpt4all-full-multi-turn"
 
 # dataset
 streaming: false
@@ -16,7 +16,7 @@ lr: 5.0e-5
 eval_every: 800
 eval_steps: 100
 save_every: 800
-output_dir: # CHANGE
+output_dir: "ckpts/gpt4all-full-multi"
 checkpoint: null
 lora: false
 warmup_steps: 100

configs/train/finetune_gptj.yaml

@@ -2,14 +2,14 @@
 model_name: "EleutherAI/gpt-j-6B"
 tokenizer_name: "EleutherAI/gpt-j-6B"
 gradient_checkpointing: true
-save_name: # CHANGE
+save_name: "nomic-ai/gpt4all-mosaic"
 
 # dataset
 streaming: false
 num_proc: 64
-dataset_path: # CHANGE
+dataset_path: "nomic-ai/turbo-500k-multi"
 max_length: 1024
-batch_size: 32
+batch_size: 8
 
 # train dynamics
 lr: 2.0e-5
@@ -18,8 +18,8 @@ weight_decay: 0.0
 eval_every: 500
 eval_steps: 105
 save_every: 500
-log_grads_every: 100
-output_dir: # CHANGE
+log_grads_every: 500
+output_dir: "ckpts/gpt4all-gptj-multinode"
 checkpoint: null
 lora: false
 warmup_steps: 500
@@ -27,7 +27,7 @@ num_epochs: 2
 
 # logging
 wandb: true
-wandb_entity: # CHANGE
-wandb_project_name: # CHANGE
+wandb_entity: vicuna
+wandb_project_name: vicuna
 seed: 42
 

configs/train/finetune_gptj_lora.yaml

@@ -2,14 +2,14 @@
 model_name: "EleutherAI/gpt-j-6b"
 tokenizer_name: "EleutherAI/gpt-j-6b"
 gradient_checkpointing: false
-save_name: # CHANGE
+save_name: "nomic-ai/gpt4all-mosaic"
 
 # dataset
 streaming: false
 num_proc: 64
-dataset_path: # CHANGE
+dataset_path: "nomic-ai/turbo-500k-multi"
 max_length: 1024
-batch_size: 1 
+batch_size: 2
 
 # train dynamics
 lr: 2.0e-5
@@ -19,7 +19,7 @@ eval_every: 500
 eval_steps: 105
 save_every: 500
 log_grads_every: 500
-output_dir: # CHANGE
+output_dir: "ckpts/gpt4all-gptj-multinode"
 checkpoint: null
 lora: true
 warmup_steps: 500
@@ -27,7 +27,7 @@ num_epochs: 2
 
 # logging
 wandb: true
-wandb_entity: # CHANGE
-wandb_project_name: # CHANGE
+wandb_entity: zanussbaum
+wandb_project_name: mosaic
 seed: 42
 

configs/train/finetune_gptjr.yaml

@@ -1,43 +0,0 @@
-# model/tokenizer
-model_name: "EleutherAI/gpt-j-6b"
-tokenizer_name: "EleutherAI/gpt-j-6b"
-version: null
-gradient_checkpointing: true
-save_name: "nomic-ai/gpt-jr-decay-alpha"
-push_to_hub: false
-encoder_dim: 384
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "/home/paperspace/gpt4all/gpt4all/index/squad_supplemented_train" 
-max_length: 1024
-batch_size: 8
-pct_test: 0.05
-q_column: "question"
-a_column: "answers"
-encoder_column: "neighbor_embeddings"
-
-
-# train dynamics
-lr: 1.0e-4
-min_lr: 0 
-weight_decay: 0.0
-eval_every: 50 
-save_every: -1
-log_grads_every: 100
-log_lr_every: 10
-output_dir: "ckpts/decay_alpha"
-checkpoint: null
-lora: false
-warmup_steps: 500
-num_epochs: 5 
-debug: false
-scheduler: false
-
-# logging
-wandb: true
-wandb_entity: gpt4all
-wandb_project_name: retrieval
-seed: 42
-

configs/train/finetune_lora.yaml

@@ -2,12 +2,12 @@
 model_name: # update
 tokenizer_name: # update
 gradient_checkpointing: false
-save_name: # CHANGE
+save_name: "nomic-ai/gpt4all-lora-llama"
 
 # dataset
 streaming: false
 num_proc: 64
-dataset_path: # CHANGE
+dataset_path: "nomic-ai/turbo-500k-multi"
 max_length: 1024
 batch_size: 4
 
@@ -18,7 +18,7 @@ weight_decay: 0.0
 eval_every: 2000
 eval_steps: 100
 save_every: 2000
-output_dir: # CHANGE
+output_dir: "ckpts/gpt4all-lora-multi"
 checkpoint: null
 lora: true
 warmup_steps: 100

configs/train/finetune_memory.yaml

@@ -1,46 +0,0 @@
-# model/tokenizer
-model_name: "EleutherAI/pythia-1b"
-tokenizer_name: "EleutherAI/pythia-1b"
-version: null
-gradient_checkpointing: true
-save_name: "nomic-ai/lethe"
-push_to_hub: false
-memory_attn_layer: 12
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "/home/paperspace/gpt4all/gpt4all/inference/synth_data_combined_174"
-max_length: 1024
-batch_size: 32
-pct_test: 0.05
-q_column: "question"
-a_column: "answer"
-context_column: "text"
-num_memories_per_index: 2000000
-num_neighbors_to_retrieve: 2
-num_neighbors_to_store: 1 
-mem_chunk_size: 64
-
-# train dynamics
-lr: 1.0e-5
-min_lr: 0 
-weight_decay: 0.0
-eval_every: 100 
-save_every: 100
-log_grads_every: 100
-log_lr_every: 10
-output_dir: "ckpts/mem_attn_no_cosine_sim"
-checkpoint: null
-lora: false
-warmup_steps: 200
-num_epochs: 5 
-debug: false
-scheduler: false
-
-# logging
-wandb: true
-wandb_entity: gpt4all
-wandb_project_name: mem_attn
-seed: 42
-

configs/train/finetune_pythiaseek.yaml

@@ -1,46 +0,0 @@
-# model/tokenizer
-model_name: "EleutherAI/pythia-1b"
-tokenizer_name: "EleutherAI/pythia-1b"
-version: null
-gradient_checkpointing: true
-save_name: "nomic-ai/pythiaseek-large-bs"
-push_to_hub: false
-encoder_dim: 384
-learnable_alpha: true
-cross_attn_layer: 12
-freeze_pretrained: false
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "/home/paperspace/gpt4all/gpt4all/index/squad_supplemented_train" 
-max_length: 1024
-batch_size: 64
-pct_test: 0.05
-q_column: "question"
-a_column: "answers"
-encoder_column: "neighbor_embeddings"
-
-
-# train dynamics
-lr: 1.0e-4
-min_lr: 0 
-weight_decay: 0.0
-eval_every: 100 
-save_every: -1
-log_grads_every: 100
-log_lr_every: 10
-output_dir: "ckpts/learnable_alpha_6b"
-checkpoint: null
-lora: false
-warmup_steps: 500
-num_epochs: 5 
-debug: false
-scheduler: false
-
-# logging
-wandb: true
-wandb_entity: gpt4all
-wandb_project_name: retrieval
-seed: 42
-

configs/train/pretrain_minipile.yaml

@@ -1,43 +0,0 @@
-# model/tokenizer
-model_name: "EleutherAI/pythia-1b"
-tokenizer_name: "EleutherAI/pythia-1b"
-version: null
-gradient_checkpointing: true
-save_name: "nomic-ai/minipille"
-push_to_hub: false
-memory_attn_layer: 12
-
-# dataset
-streaming: false
-num_proc: 64
-dataset_path: "JeanKaddour/minipile" 
-max_length: 2048
-batch_size: 64
-pct_test: 0.05
-num_memories_per_index: 5000000
-mem_chunk_size: 512
-num_chunks: 10
-num_neighbors_to_retrieve: 32
-
-# train dynamics
-lr: 1.0e-4
-min_lr: 0 
-weight_decay: 0.0
-eval_every: 100 
-save_every: -1
-log_grads_every: 100
-log_lr_every: 10
-output_dir: "ckpts/minipile"
-checkpoint: null
-lora: false
-warmup_steps: 500
-num_epochs: 5 
-debug: false
-scheduler: false
-
-# logging
-wandb: true
-wandb_entity: gpt4all
-wandb_project_name: minipile
-seed: 42
-

create_hostname.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+export WORKER_IP=$1
+N_GPUS=8
+# create dir if doesn't exist
+sudo mkdir -p /job
+printf "localhost slots=$N_GPUS\n$WORKER_IP slots=$N_GPUS" | sudo tee /job/hostfile
+echo /job/hostfile

data.py

@@ -1,17 +1,68 @@
 import glob
 import torch
-from datasets import load_dataset
+from datasets import load_dataset, concatenate_datasets
 import os
-import hnswlib
 from torch.utils.data import DataLoader
 from transformers import DefaultDataCollator
-from .preprocess import tokenize_inputs
+
+
+
+def tokenize_inputs(config, tokenizer, examples):
+    max_length = config["max_length"]
+
+    # hacky backward compatible
+    different_eos = tokenizer.eos_token != "</s>"
+    out = {"labels": [], "input_ids": []}
+    for prompt, response in zip(examples["prompt"], examples["response"]):
+        if different_eos:
+            if response.count("</s>") > 0:
+                response = response.replace("</s>", tokenizer.eos_token)
+
+        prompt_len = len(tokenizer(prompt + "\n", return_tensors="pt")["input_ids"][0])
+
+        # hack if our prompt is super long
+        # we need to include some labels so we arbitrarily trunacate at max_length // 2
+        # if the length is too long
+        if prompt_len >= max_length // 2:
+            # if prompt is too long, truncate
+            # but make sure to truncate to at max 1024 tokens
+            new_len = min(max_length // 2, len(prompt) // 2)
+            prompt = prompt[:new_len]
+            # get new prompt length
+            prompt_len = tokenizer(prompt + "\n", return_tensors="pt", max_length=max_length // 2, truncation=True).input_ids.ne(tokenizer.pad_token_id).sum().item()
+
+        assert prompt_len <= max_length // 2, f"prompt length {prompt_len} exceeds max length {max_length}"
+
+        input_tokens = tokenizer(prompt + "\n" + response + tokenizer.eos_token,
+                                 truncation=True, max_length=max_length, return_tensors="pt")["input_ids"].squeeze()
+
+        labels = input_tokens.clone()
+        labels[:prompt_len] = -100
+        if len(labels) < max_length:
+            # pad to max_length with -100
+            labels = torch.cat([labels, torch.full((max_length - len(labels),), -100)])
+
+        assert (labels == -100).sum() < len(labels), f"Labels are all -100, something wrong. prompt length {prompt_len} exceeds max length {max_length}" 
+        
+        if (labels == -100).sum() == len(labels) - 1:
+            print(prompt)
+            print(response)
+            raise
+
+        input_tokens = tokenizer.pad({"input_ids": input_tokens}, padding="max_length", max_length=max_length)["input_ids"]
+        out["labels"].append(labels)
+        out["input_ids"].append(input_tokens)
+
+    out = {k: torch.stack(v) if isinstance(v, list) else v for k, v in out.items()}
+
+    return out
 
 
 def load_data(config, tokenizer):
     dataset_path = config["dataset_path"]
 
     if os.path.exists(dataset_path):
+        # check if path is a directory
         if os.path.isdir(dataset_path):
             files = glob.glob(os.path.join(dataset_path, "*_clean.jsonl"))
         else:
@@ -35,13 +86,13 @@ def load_data(config, tokenizer):
 
     # tokenize inputs and return labels and attention mask
     train_dataset = train_dataset.map(
-        lambda ele: tokenize_inputs(config, tokenizer, ele, "prompt", "response"),
+        lambda ele: tokenize_inputs(config, tokenizer, ele),
         batched=True,
         remove_columns=["source", "prompt"],
         **kwargs
     )
     val_dataset = val_dataset.map(
-        lambda ele: tokenize_inputs(config, tokenizer, ele, "prompt", "response"),
+        lambda ele: tokenize_inputs(config, tokenizer, ele), 
         batched=True,
         remove_columns=["source", "prompt"],
         **kwargs
@@ -66,7 +117,6 @@ def load_data(config, tokenizer):
 
     return train_dataloader, val_dataloader
 
-
     
 def load_data_for_inference(config, tokenizer):
     dataset_path = config["dataset_path"]
@@ -103,12 +153,12 @@ def load_data_for_inference(config, tokenizer):
 
     # tokenize inputs and return labels and attention mask
     train_dataset = train_dataset.map(
-        lambda ele: tokenize_inputs(config, tokenizer, ele, "prompt", "response"),
+        lambda ele: tokenize_inputs(config, tokenizer, ele),
         batched=True,
         **kwargs
     )
     val_dataset = val_dataset.map(
-        lambda ele: tokenize_inputs(config, tokenizer, ele, "prompt", "response"), 
+        lambda ele: tokenize_inputs(config, tokenizer, ele), 
         batched=True,
         **kwargs
     )

eval_figures.py

@@ -6,20 +6,18 @@ from matplotlib import pyplot as plt
 plt.figure()
 for fpath in glob.glob('./eval_data/*.pkl'):
     parts = fpath.split('__')
-    model_name = "-".join(fpath.replace(".pkl", "").split("_")[2:])
+    model_name = parts[1].replace('model-', '').replace('.pkl', '')
+    lora_name = parts[2].replace('lora-', '').replace('.pkl', '')
     with open(fpath, 'rb') as f:
         data = pickle.load(f)
         perplexities = data['perplexities']
         perplexities = np.nan_to_num(perplexities, 100)
         perplexities = np.clip(perplexities, 0, 100)
-        if 'alpaca' not in fpath:
-            identifier = model_name = "-".join(fpath.replace(".pkl", "").split("eval__model-")[1:]) 
-            label = 'GPT4all-'
-            label += identifier
-            
+        if 'nomic' in fpath:
+            label = 'GPT4all-lora'
         else:
             label = 'alpaca-lora'
-        plt.hist(perplexities, label=label, alpha=.5, bins=50)
+        plt.hist(perplexities, label=label, alpha=.5)
 
 plt.xlabel('Perplexity')
 plt.ylabel('Frequency')

eval_self_instruct.py

@@ -3,7 +3,7 @@ import torch
 import pickle
 import numpy as np
 from tqdm import tqdm
-from gpt4all.utils.read import read_config
+from read import read_config
 from argparse import ArgumentParser
 from peft import PeftModelForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -49,6 +49,28 @@ def eval_example(model, tokenizer, example, config):
     input = tokenizer(prompt, return_tensors="pt")
     input = {k: v.to(model.device) for k, v in input.items()}
 
+    continuations = []
+    tokenized_continuations = []
+    trajectories = []
+    for i in range(1):
+        with torch.no_grad():
+            outputs = model.generate(input_ids=input['input_ids'],
+                                     max_new_tokens=config["max_new_tokens"],
+                                     min_new_tokens=5,
+                                     temperature=config["temperature"],
+                                     repetition_penalty=1.0,
+                                     do_sample=True)
+            decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+            y = model(input_ids=outputs)
+        trajectory = y.hidden_states[0].detach().cpu().numpy()[0]
+        trajectory = trajectory / np.linalg.norm(trajectory, axis=1, keepdims=True)
+        trajectory = np.cumsum(trajectory, axis=0) / np.arange(1, trajectory.shape[0]+1).reshape(-1, 1)
+
+        trajectories.append(trajectory)
+        continuations.append(decoded)
+        tokenized_continuations.append(tokenizer.tokenize(decoded))
+
     #compute the ground truth perplexity
     gt_input = tokenizer(gt, return_tensors="pt")
     gt_input = {k: v.to(model.device) for k, v in gt_input.items()}
@@ -79,23 +101,30 @@ def eval_example(model, tokenizer, example, config):
 
     print(prompt)
     print(80*'-')
-   
+    for continuation in continuations:
+        print(continuation)
+        print(80*'-')
 
-    return ppl
+    return ppl, trajectories, continuations, tokenized_continuations
 
 def do_eval(config):
     eval_data = read_jsonl_file('eval_data/user_oriented_instructions.jsonl')
     model, tokenizer = setup_model(config)
+    all_trajectories = []
     all_perplexities = []
+    all_continuations = []
+    all_tokenized_continuations = []
     for example in tqdm(eval_data):
-        gt_perplexity = eval_example(model, tokenizer, example, config)
+        gt_perplexity, trajectories, continuations, tokenized_continuations = eval_example(model, tokenizer, example, config)
+        all_trajectories.append(trajectories)
         all_perplexities.append(gt_perplexity)
+        all_continuations.append(continuations)
 
-        
-    name = f"eval_data/eval__model-{config['model_name'].replace('/', '_')}{'__lora-' + config['lora_path'].replace('/', '_') if config['lora'] else ''}.pkl"
-
-    with open(name, 'wb') as f:
-        r = {'perplexities': all_perplexities}
+    with open('eval_data/eval__model-{}__lora-{}.pkl'.format(config['model_name'].replace('/', '_'), config['lora_path'].replace('/', '_')), 'wb') as f:
+        r = {'trajectories': all_trajectories,
+             'perplexities': all_perplexities,
+             'continuations': all_continuations,
+             'tokenized_continuations': all_tokenized_continuations}
         pickle.dump(r, f)
 
 

generate.py

@@ -1,6 +1,6 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModelForCausalLM
-from gpt4all.utils.read import read_config
+from read import read_config
 from argparse import ArgumentParser
 import torch
 import time

gpt4all/data/__init__.py

@@ -1,2 +0,0 @@
-from .instruction_tuning_dataloader import *
-from .retrieval_dataloader import *

gpt4all/data/preprocess.py

@@ -1,51 +0,0 @@
-import torch
-
-def tokenize_inputs(config, tokenizer, examples, input_col, target_col):
-    max_length = config["max_length"]
-
-    # hacky backward compatible
-    different_eos = tokenizer.eos_token != "</s>"
-    out = {"labels": [], "input_ids": []}
-    for prompt, response in zip(examples[input_col], examples[target_col]):
-        if different_eos:
-            if response.count("</s> \n") > 0:
-                response = response.replace("</s> \n", f"{tokenizer.eos_token} \n") 
-
-        prompt_len = len(tokenizer(prompt + "\n", return_tensors="pt")["input_ids"][0])
-
-        # hack if our prompt is super long
-        # we need to include some labels so we arbitrarily trunacate at max_length // 2
-        # if the length is too long
-        if prompt_len >= max_length // 2:
-            # if prompt is too long, truncate
-            # but make sure to truncate to at max 1024 tokens
-            new_len = min(max_length // 2, len(prompt) // 2)
-            prompt = prompt[:new_len]
-            # get new prompt length
-            prompt_len = tokenizer(prompt + "\n", return_tensors="pt", max_length=max_length // 2, truncation=True).input_ids.ne(tokenizer.pad_token_id).sum().item()
-
-        assert prompt_len <= max_length // 2, f"prompt length {prompt_len} exceeds max length {max_length}"
-
-        input_tokens = tokenizer(prompt + "\n" + response + tokenizer.eos_token,
-                                 truncation=True, max_length=max_length, return_tensors="pt")["input_ids"].squeeze()
-
-        labels = input_tokens.clone()
-        labels[:prompt_len] = -100
-        if len(labels) < max_length:
-            # pad to max_length with -100
-            labels = torch.cat([labels, torch.full((max_length - len(labels),), -100)])
-
-        assert (labels == -100).sum() < len(labels), f"Labels are all -100, something wrong. prompt length {prompt_len} exceeds max length {max_length}" 
-        
-        if (labels == -100).sum() == len(labels) - 1:
-            print(prompt)
-            print(response)
-            raise
-
-        input_tokens = tokenizer.pad({"input_ids": input_tokens}, padding="max_length", max_length=max_length)["input_ids"]
-        out["labels"].append(labels)
-        out["input_ids"].append(input_tokens)
-
-    out = {k: torch.stack(v) if isinstance(v, list) else v for k, v in out.items()}
-
-    return out

gpt4all/data/retrieval_dataloader.py

@@ -1,203 +0,0 @@
-from datasets import load_dataset, Dataset
-import os
-from torch.utils.data import DataLoader
-from .preprocess import tokenize_inputs
-from transformers import DefaultDataCollator
-
-
-def load_retrieval_augmented_data(config, tokenizer, split="train", split_dataset=True):
-    dataset_path = config["dataset_path"]
-
-    if os.path.exists(dataset_path):
-        dataset = Dataset.load_from_disk(dataset_path)
-    else:
-        dataset = load_dataset(dataset_path, split=split)
-
-
-    question_col = config["q_column"]
-    answer_col = config["a_column"]
-    encoder_column = config["encoder_column"]
-
-    if config["streaming"] is False:
-        kwargs = {"num_proc": config["num_proc"]}
-    else:
-        kwargs = {}
-
-    # strip any unneccessary whitespace
-    # there's one question that's includes a ton of whitespace
-    dataset = dataset.map(lambda ele: {question_col: [q.strip() for q in ele[question_col]]}, batched=True)
-    # in squad, the data is formatted where each ele in answers is a dict where the key text holds
-    # a list of the answer
-    dataset = dataset.map(lambda ele: {answer_col: [t["text"][0] for t in ele[answer_col]]}, batched=True)
-
-    dataset = dataset.map(
-        lambda ele: tokenize_inputs(config, tokenizer, ele, question_col, answer_col),
-        batched=True,
-        **kwargs
-    )
-
-    # tokenize inputs + labels in teacher-force format
-    # rename encoder hidden states if not already called that
-    if encoder_column != "encoder_hidden_states":
-        dataset = dataset.rename_column(encoder_column, "encoder_hidden_states")
-
-    columns_to_keep = ["input_ids", "labels", "encoder_hidden_states"]
-
-    col_names_to_rm = [col for col in dataset.column_names if col not in columns_to_keep]
-    dataset = dataset.remove_columns(col_names_to_rm)
-
-    if split_dataset:
-        dataset = dataset.train_test_split(test_size=config["pct_test"], seed=config["seed"])
-        train_dataset, val_dataset = dataset["train"], dataset["test"]
-
-        train_dataloader = DataLoader(
-            train_dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        val_dataloader = DataLoader(
-            val_dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        return train_dataloader, val_dataloader
-
-    else:
-        dataloader = DataLoader(
-            dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        return dataloader
-
-
-def load_memory_augmented_data(config, tokenizer, split="train", split_dataset=True):
-    dataset_path = config["dataset_path"]
-
-    if os.path.exists(dataset_path):
-        dataset = Dataset.load_from_disk(dataset_path)
-    else:
-        dataset = load_dataset(dataset_path, split=split)
-
-
-    question_col = config["q_column"]
-    answer_col = config["a_column"]
-    context_col = config["context_column"]
-
-    if config["streaming"] is False:
-        kwargs = {"num_proc": config["num_proc"]}
-    else:
-        kwargs = {}
-
-    # strip any unneccessary whitespace
-    # there's one question that's includes a ton of whitespace
-    dataset = dataset.map(lambda ele: {question_col: [q.strip() for q in ele[question_col]]}, batched=True)
-    # in squad, the data is formatted where each ele in answers is a dict where the key text holds
-    # a list of the answer
-    dataset = dataset.map(lambda ele: {answer_col: [t.strip() for t in ele[answer_col]]}, batched=True)
-    # dataset = dataset.map(lambda ele: {answer_col: [t["text"][0] for t in ele[answer_col]]}, batched=True)
-
-    dataset = dataset.map(
-        lambda ele: tokenize_inputs(config, tokenizer, ele, question_col, answer_col),
-        batched=True,
-        **kwargs
-    )
-
-    # tokenize contexts for each example
-    dataset = dataset.map(
-        lambda ele: {"retrieved_context": tokenizer([ele[context_col]], 
-                                                    return_tensors="pt", 
-                                                    padding="max_length", 
-                                                    truncation=True)["input_ids"]},
-        **kwargs
-    )
-
-    columns_to_keep = ["id", "input_ids", "labels", "retrieved_context"]
-
-    col_names_to_rm = [col for col in dataset.column_names if col not in columns_to_keep]
-    dataset = dataset.remove_columns(col_names_to_rm)
-
-    if split_dataset:
-        dataset = dataset.train_test_split(test_size=config["pct_test"], seed=config["seed"])
-        train_dataset, val_dataset = dataset["train"], dataset["test"]
-
-        train_dataloader = DataLoader(
-            train_dataset.remove_columns("id"),
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        val_dataloader = DataLoader(
-            val_dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        return train_dataloader, val_dataloader 
-
-    else:
-        dataloader = DataLoader(
-            dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        return dataloader
-
-
-def load_memory_pretraining_data(config, tokenizer, split="train", split_dataset=True):
-    dataset_path = config["dataset_path"]
-
-    if os.path.exists(dataset_path):
-        dataset = Dataset.load_from_disk(dataset_path)
-    else:
-        dataset = load_dataset(dataset_path, split=split)
-
-    if config["streaming"] is False:
-        kwargs = {"num_proc": config["num_proc"]}
-    else:
-        kwargs = {}
-        
-    # e.g. 512 * 10 = 5120 sequence length split up
-    max_length = config["mem_chunk_size"] * config["num_chunks"]
-    dataset = dataset.map(lambda ele: tokenizer(ele["text"], padding="max_length", truncation=True, max_length=max_length),
-                          batched=True, **kwargs)
-
-    dataset = dataset.map(lambda x: {"labels": x["input_ids"]}, batched=True, **kwargs)
-
-
-    columns_to_keep = ["input_ids", "labels", "attention_mask"]
-
-    col_names_to_rm = [col for col in dataset.column_names if col not in columns_to_keep]
-    dataset = dataset.remove_columns(col_names_to_rm)
-
-    # we can shuffle since the docs are in one row not split across rows
-    if split_dataset:
-        dataset = dataset.train_test_split(test_size=config["pct_test"], seed=config["seed"])
-        train_dataset, val_dataset = dataset["train"], dataset["test"]
-
-        train_dataloader = DataLoader(
-            train_dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        val_dataloader = DataLoader(
-            val_dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        return train_dataloader, val_dataloader
-
-    else:
-        dataloader = DataLoader(
-            dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-        )
-
-        return dataloader

gpt4all/eval/eval_squad.py

@@ -1,59 +0,0 @@
-import torch
-from gpt4all.models import GPTJRForCausalLM, PythiaSeekForCausalLM
-from gpt4all.data.retrieval_dataloader import load_retrieval_augmented_data
-from gpt4all.train.metrics import f1_score, exact_match_score
-from gpt4all.utils.read import read_config
-from transformers import AutoTokenizer
-from argparse import ArgumentParser
-from tqdm import tqdm
-
-parser = ArgumentParser()
-parser.add_argument("--config", type=str, default="config.yaml")
-
-args = parser.parse_args()
-
-config = read_config(args.config)
-
-tokenizer = AutoTokenizer.from_pretrained(config["tokenizer_name"])
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-dataloader = load_retrieval_augmented_data(config, tokenizer, split_dataset=False)
-
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = PythiaSeekForCausalLM.from_pretrained(config["model_name"], use_cache=False)
-model.to(device)
-model.eval()
-
-# Evaluate the model on the SQUAD dataset
-f1s = []
-exact_matches = []
-with torch.no_grad():
-    for batch in tqdm(dataloader):
-        outputs = model(input_ids=batch["input_ids"].to(device),
-                        labels=batch["labels"].to(device),
-                        encoder_hidden_states=batch["encoder_hidden_states"].to(device))
-
-        labels = batch["labels"]
-        mask = labels == -100
-
-        # since it's causal we predict the next token
-        predicted_tokens = outputs.logits.argmax(dim=-1)[:, :-1]
-        predicted_tokens[mask[:, 1:]] = tokenizer.pad_token_id
-        predicted = tokenizer.batch_decode(predicted_tokens, skip_special_tokens=True)
-
-        labels[mask] = tokenizer.pad_token_id
-        ground_truth = tokenizer.batch_decode(labels, skip_special_tokens=True)
-        print(f"Predicted: {predicted}")
-        print(f"Ground truth: {ground_truth}")
-
-        f1 = f1_score(predicted, ground_truth)
-        exact_match = exact_match_score(predicted, ground_truth)
-
-        f1s.extend(f1)
-        exact_matches.extend(exact_match)
-
-
-print(torch.tensor(f1s).mean())
-print(torch.tensor(exact_matches).to(torch.float32).mean())

gpt4all/eval/eval_squad_atlas_map.py

@@ -1,116 +0,0 @@
-import torch
-import torch.nn.functional as F
-from gpt4all.models import LetheForCausalLM
-from gpt4all.models.lethe.modeling_lethe import MemoryIndex
-from gpt4all.data.retrieval_dataloader import load_memory_augmented_data
-from gpt4all.train.metrics import f1_score, exact_match_score
-from gpt4all.utils.read import read_config
-from transformers import AutoTokenizer, AutoConfig
-from argparse import ArgumentParser
-from tqdm import tqdm
-from nomic import atlas
-from datasets import load_from_disk
-
-
-def calc_loss_per_item(logits, labels):
-    lm_logits = logits[:, :-1, :].contiguous()
-    lm_labels = labels[:, 1:].contiguous()
-    loss = F.cross_entropy(lm_logits.view(-1, lm_logits.size(-1)), lm_labels.view(-1), reduction="none")
-    loss = loss.reshape(labels.shape[0], -1).mean(dim=-1)
-
-    # return tensor of shape (B,) where B is the batch size
-    return loss.cpu().tolist()
-
-    
-def greedy_search(input_ids, model, tokenizer, max_new_tokens=100):
-    num_new_tokens = 0
-    with torch.no_grad():
-        while True:
-            if num_new_tokens >= max_new_tokens:
-                break
-            outputs = model(input_ids, save_kv=False)
-
-            new_tokens = torch.argmax(outputs.logits[:, -1, :], dim=-1)
-
-            input_ids = torch.cat([input_ids, new_tokens.unsqueeze(1)], dim=-1)
-            num_new_tokens += 1
-
-            if torch.equal(input_ids[0, -1].cpu(), torch.tensor(tokenizer.eos_token_id)):
-                break
-        
-    print(tokenizer.batch_decode(input_ids, skip_special_tokens=True))
-
-    return input_ids
-    
-
-parser = ArgumentParser()
-parser.add_argument("--config", type=str, default="config.yaml")
-
-args = parser.parse_args()
-
-config = read_config(args.config)
-
-tokenizer = AutoTokenizer.from_pretrained(config["tokenizer_name"], model_max_length=config["max_length"])
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-dataloader = load_memory_augmented_data(config, tokenizer, split_dataset=False)
-
-dataset = load_from_disk(config["dataset_path"])
-
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model_config = AutoConfig.from_pretrained(config["model_name"])
-
-head_size = model_config.hidden_size // model_config.num_attention_heads
-index = MemoryIndex(head_size,
-            config["num_memories_per_index"],
-            model_config.num_attention_heads
-)
-model = LetheForCausalLM.from_pretrained(config["model_name"], 
-                                    revision=config['version'] if 'version' in config else None,
-                                    memory_attn_layer=config["memory_attn_layer"],
-                                    num_neighbors_to_retrieve=config["num_neighbors_to_retrieve"],
-                                    index=index,
-                                ).to(device)
-model.eval()
-
-# Evaluate the model on the SQUAD dataset
-losses = []
-with torch.no_grad():
-    for i, batch in enumerate(tqdm(dataloader)):
-        memories = batch["retrieved_context"]
-        memories = memories[:, :config["num_neighbors_to_store"], :]
-        memories = memories.reshape(-1, memories.shape[-1])
-
-        # need to set to eval so we don't do mem attn as it's slow
-        with torch.no_grad():
-            for chunk_start in range(0, memories.shape[0], config["mem_chunk_size"]):
-                chunk_end = min(memories.shape[0], chunk_start + config["mem_chunk_size"])
-                mem_chunk = memories[chunk_start:chunk_end]
-                model(input_ids=mem_chunk.to(device))
-
-        del memories
-        torch.cuda.empty_cache()
-        qa_inputs = batch["input_ids"]
-        qa_labels = batch["labels"]
-        for i in range(qa_inputs.shape[0]):
-            inputs = qa_inputs[i].to(device)
-            labels = qa_labels[i].to(device)
-            cutoff = torch.argmax((labels != -100).type(torch.float32))
-            greedy_search(inputs[:cutoff.item()].unsqueeze(0).to(device), model, tokenizer)
-            print(tokenizer.decode(inputs, skip_special_tokens=True))
-
-
-        # batch_loss = calc_loss_per_item(outputs.logits, qa_labels.to(device))
-        # losses.extend(batch_loss)
-        index.reset()
-
-        
-        
-dataset = dataset.add_column("loss", losses)
-
-dataset.save_to_disk("eval_squad_atlas_map")
-
-
-        

gpt4all/index/README.md

@@ -1,34 +0,0 @@
-# How to Tokenize and Embed
-
-Split text into chunks
-```
-python tokenize_texts.py
-```
-
-Embbed Texts
-
-```
-torchrun --master_port=29085 --nproc-per-node 8  embed_texts.py --ds_path=tokenized --batch_size=2048
-```
-
-
-Combine Embeddings and Build Index
-```
-python build_index.py --ds_path=wiki_sample_tokenized --embed_folder=wiki_sample_embedded
-```
-
-To use the Index
-
-```
-import hnswlib
-
-index = hnswlib.Index(space='l2', dim=384)
-index.load_index(<path to index>)
-```
-
-
-Prep index for train
-
-```
-CUDA_VISIBLE_DEVICES=7 torchrun --master_port=29086 --nproc-per-node 1 prep_index_for_train.py --config=../../configs/train/finetune_gptjr.yaml
-```

gpt4all/index/build_index.py

@@ -1,102 +0,0 @@
-import os
-from datasets import Dataset, concatenate_datasets
-import glob
-from argparse import ArgumentParser
-import hnswlib
-import pyarrow as pa
-import pyarrow.compute as pc
-from tqdm import tqdm
-
-
-def parse_args():
-    parser = ArgumentParser()
-    parser.add_argument("--ds_path", type=str, required=True)
-    parser.add_argument("--embed_folder", type=str, required=True)
-    parser.add_argument("--index_path", type=str, default="wiki-index")
-
-    return parser.parse_args()
-
-
-def concat_embedded(folder):
-    files = glob.glob(f"{folder}/*")
-
-    all_embeddings = []
-    for file in files:
-        all_embeddings.append(Dataset.load_from_disk(file))
-
-    all_embeddings = concatenate_datasets(all_embeddings)
-
-    return all_embeddings
-
-
-def join(original_ds, embedded_ds):
-    embedded_ds = embedded_ds.add_column("index", range(len(embedded_ds)))
-    embed_table = embedded_ds.data.table
-
-    seen = set()
-    indices = []
-    for i, id in enumerate(original_ds["id"]):
-        if id not in seen:
-            indices.append(i)
-            seen.add(id)
-
-    mask = pc.is_in(embed_table["index"], value_set=pa.array(indices, pa.int32()))
-    filtered_table = embed_table.filter(mask)
-
-    import pdb; pdb.set_trace()
-    # sort to make sure we're adding in right order
-    filtered_table = filtered_table.sort_by("id")
-
-    original_table = original_ds.data.table
-    original_table = original_table.sort_by("id")
-
-    original_table = original_table.append_column(
-        "embedding", filtered_table["embedding"]
-    )
-    # there's definitely a better way to do this but
-    # Dataset(original_table) throws `KeyError: 'embedding'`
-    joined = Dataset.from_dict(original_table.to_pydict())
-
-    return joined
-
-
-def build_index(orig_path, embed_folder_path, index_path):
-    if not os.path.exists(orig_path + "_embedded_with_text"):
-        # TODO: this doesn't work for large datasets!
-        # just convert to pandas and do this manually
-        ds = Dataset.load_from_disk(orig_path)
-        embed_ds = concat_embedded(embed_folder_path)
-        print("Concatenated embeddings")
-        print(f"Length: {len(ds)}")
-        print(f"Length: {len(embed_ds)}")
-        ds = join(ds, embed_ds)
-        ds = ds.add_column("index", range(len(ds)))
-        print("Saving to disk")
-        ds.save_to_disk(f"{orig_path}_embedded_with_text")
-    else:
-        ds = Dataset.load_from_disk(orig_path + "_embedded_with_text")
-
-    print(f"Length of ds: {len(ds)}")
-
-    print("Building index")
-    index = hnswlib.Index(space="cosine", dim=384)
-    # not sure what we should set M and ef_construction to
-    index.init_index(max_elements=len(ds), M=64, ef_construction=200)
-    print("Adding items")
-    chunk_size = 50_000
-    num_chunks = len(ds) // chunk_size
-    progbar = tqdm(total=num_chunks)
-    start = 0
-    while start < len(ds):
-        chunk = ds[start:start + chunk_size]
-        index.add_items(chunk["embedding"], chunk["id"], num_threads=64)
-        progbar.update(1)
-        start += chunk_size
-
-    print("Saving index")
-    index.save_index(index_path + ".bin")
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    build_index(args.ds_path, args.embed_folder, args.index_path)

gpt4all/index/embed.py

@@ -1,70 +0,0 @@
-import torch
-from transformers import AutoTokenizer, AutoModel
-import torch.nn.functional as F
-
-
-class Embedder:
-    def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.embedder = AutoModel.from_pretrained(model_name)
-        # hack
-        self.offset = self.tokenizer.model_max_length // 2
-
-    def _mean_pool(self, model_output, attention_mask):
-        token_embeddings = model_output[
-            0
-        ]  # First element of model_output contains all token embeddings
-        input_mask_expanded = (
-            attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-        )
-        sentence_embeddings = torch.sum(
-            token_embeddings * input_mask_expanded, 1
-        ) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-        return F.normalize(sentence_embeddings, p=2, dim=1)
-
-    def chunk_text(self, text):
-        tokenized_text = {"input_ids": [], "attention_mask": []}
-        tokenized = self.tokenizer(text)
-        tokenized_len = len(tokenized["input_ids"])
-        max_len = self.tokenizer.model_max_length
-        if tokenized_len > max_len:
-            start = 0
-            while start < tokenized_len:
-                tokenized_text["input_ids"].append(
-                    tokenized["input_ids"][start : start + max_len]
-                )
-                tokenized_text["attention_mask"].append(
-                    tokenized["attention_mask"][start : start + max_len]
-                )
-                # this could probably be done better
-                start += self.offset
-
-        else:
-            tokenized_text["input_ids"].append(tokenized["input_ids"])
-            tokenized_text["attention_mask"].append(tokenized["attention_mask"])
-
-        return tokenized_text
-
-    def tokenize(self, text):
-        return self.tokenizer(text, return_tensors="pt", truncation=True, padding="max_length")
-
-    def __call__(self, batch):
-        if isinstance(batch, str):
-            tokenized = self.tokenizer(batch, return_tensors="pt", truncation=True)
-            return self._mean_pool(
-                self.embedder(
-                    input_ids=tokenized["input_ids"],
-                    attention_mask=tokenized["attention_mask"],
-                ),
-                tokenized["attention_mask"],
-            )
-        else:
-            outputs = self.embedder(
-                input_ids=batch["input_ids"], attention_mask=batch["attention_mask"]
-            )
-            embedding = self._mean_pool(outputs, batch["attention_mask"])
-
-            return {"id": batch["id"], "embedding": embedding}
-
-    def to(self, device):
-        self.embedder.to(device)

gpt4all/index/embed_texts.py

@@ -1,116 +0,0 @@
-import os
-import torch.distributed as dist
-from argparse import ArgumentParser
-from datasets import Dataset
-from gpt4all.index.embed import Embedder
-from gpt4all.utils.distributed_utils import rank0_print
-from torch.utils.data import DataLoader, DistributedSampler
-from transformers.trainer_pt_utils import nested_numpify
-from transformers import BatchEncoding
-from tqdm import tqdm
-import numpy as np
-import torch
-from datasets import load_dataset
-
-
-# this isn't used but keeping in case we need it in the future
-# collate and pad inputs to the right shape
-class PadCollateInputs:
-    def __init__(self, tokenizer):
-        self.tokenizer = tokenizer
-
-    def __call__(self, batch):
-        mapped_inputs = {"input_ids": [], "attention_mask": []}
-        mapped_inputs["input_ids"] = [b["tokenized_chunk"] for b in batch]
-        mapped_inputs["attention_mask"] = [b["tokenized_attn_mask"] for b in batch]
-        encoding = BatchEncoding(mapped_inputs)
-
-        padded_inputs = self.tokenizer.pad(
-            encoding, padding="max_length", return_tensors="pt"
-        )
-        padded_inputs["id"] = [b["id"] for b in batch]
-
-        return padded_inputs
-
-
-def embed_texts(ds_path, batch_size, embed_on='text', save_to_disk=False, split='train'):
-    world_size = dist.get_world_size() if dist.is_initialized() else 1
-    rank0_print(f"World size: {world_size}")
-
-    if os.path.exists(ds_path):
-        dataset = Dataset.load_from_disk(ds_path)
-    else:
-        dataset = load_dataset(ds_path, split=split)
-    rank0_print(f"Dataset size: {len(dataset)}")
-
-    model = Embedder()
-
-    if "input_ids" not in dataset.column_names:
-        dataset = dataset.map(lambda x: model.tokenize(x[embed_on]), batched=True, num_proc=64)
-
-
-    columns_to_keep = ["id", "input_ids", "attention_mask"]
-    to_remove = [e for e in dataset.column_names if not e in columns_to_keep]
-    dataset = dataset.remove_columns(to_remove)
-
-    dataset = dataset.with_format("torch")
-
-    num_processes = dist.get_world_size() if dist.is_initialized() else 1
-    local_rank = dist.get_rank() if dist.is_initialized() else 0 
-
-
-    if num_processes > 1: 
-        sampler = DistributedSampler(
-            dataset,
-            shuffle=False,
-            drop_last=False,
-            num_replicas=num_processes,
-            rank=local_rank,
-        )
-    else:
-        sampler = None
-
-    dataloader = DataLoader(
-        dataset,
-        batch_size=batch_size,
-        sampler=sampler,
-        drop_last=False,
-    )
-
-    model.to(f"cuda:{local_rank}")
-    with torch.no_grad():
-        embedded_outputs = {"id": [], "embedding": []}
-        for batch in tqdm(dataloader, disable=local_rank != 0):
-            batch["input_ids"] = batch["input_ids"].to(f"cuda:{local_rank}")
-            batch["attention_mask"] = batch["attention_mask"].to(f"cuda:{local_rank}")
-            outputs = model(batch)
-            embedded_outputs["id"].extend(batch["id"])
-            embedded_outputs["embedding"].extend(outputs["embedding"])
-
-        embedded_outputs["embedding"] = nested_numpify(embedded_outputs["embedding"])
-        embedded_outputs["id"] = np.stack(embedded_outputs["id"])
-        embedded_outputs["embedding"] = np.stack(embedded_outputs["embedding"])
-
-    ds = Dataset.from_dict(embedded_outputs)
-
-    # feeling lazy, don't want to wait for all_gather to finish
-    # will load and concat in a single process in another script
-    if save_to_disk:
-        ds.save_to_disk(f"{ds_path}_embedded/{ds_path}_embedded_rank_{local_rank}")
-    return ds
-
-
-def main():
-    dist.init_process_group("nccl")
-    parser = ArgumentParser()
-    parser.add_argument("--ds_path", type=str, default="tokenized")
-    parser.add_argument("--batch_size", type=int, default=1)
-
-    args = parser.parse_args()
-
-    embed_texts(args.ds_path, args.batch_size, save_to_disk=True)
-
-
-if __name__ == "__main__":
-    # parse arguments by reading in a config
-    main()

gpt4all/index/prep_index_for_train.py

@@ -1,94 +0,0 @@
-import os
-import hnswlib
-import numpy as np
-import pyarrow as pa
-from datasets import Dataset
-import torch.distributed as dist
-from datasets import load_dataset
-from pyarrow import compute as pc
-from argparse import ArgumentParser
-from gpt4all.utils.read import read_config
-from gpt4all.index.embed_texts import embed_texts
-from tqdm import tqdm
-
-CHUNK_SIZE = 2048
-
-def parse_args():
-    parser = ArgumentParser()
-    parser.add_argument("--config", type=str, default="config.yaml")
-    parser.add_argument("--split", type=str, default="train")
-    parser.add_argument("--k", type=int, default=5)
-
-    return parser.parse_args()
-
-    
-def prep_index():
-    args = parse_args()
-    config = read_config(args.config)
-    index = hnswlib.Index(space=config['index_space'], dim=config['index_dim'])
-    print("loading index")
-    index.load_index(config['index_path'])
-
-    # load query dataset
-    ds_path = config['dataset_path']
-
-    # load retrieval dataset
-    print("loading retrieval dataset")
-    print(config["index_database"])
-    if os.path.exists(config['index_database']):
-        retrieval_dataset = Dataset.load_from_disk(config['index_database'])
-    else:
-        retrieval_dataset = load_dataset(config['index_database'], split="train")
-
-    # vectorize queries
-    query_vector_path =  f"{ds_path}_queries_embedded/{ds_path}_embedded_{args.split}"
-    if not os.path.exists(query_vector_path):
-        print('Embedding dataset...')
-        ds = embed_texts(ds_path,
-                         config['batch_size'],
-                         embed_on=config['query_embedding_field'],
-                         save_to_disk=False,
-                         split=args.split)
-        ds.save_to_disk(query_vector_path)
-    else:
-        print('Found cached embedding dataset!')
-        ds = Dataset.load_from_disk(query_vector_path)
-
-    #build training dataset
-    train_dataset = load_dataset(ds_path, split=args.split)
-
-    # search the index for each query
-    neighbor_text_column = []
-    neighbor_ids_column = []
-    for _, chunk_start in enumerate(tqdm(range(0, len(ds), CHUNK_SIZE))):
-        chunk_end = chunk_start + CHUNK_SIZE
-        chunk = ds[chunk_start:chunk_end]
-        query_vectors = np.array(chunk['embedding'])
-        neighbor_ids, _ = index.knn_query(query_vectors, k=args.k, num_threads=-1) # neighbor ids is of shape [n_queries, n_neighbors]
-        value_set = pa.array([str(e) for e in neighbor_ids.flatten()])
-        neighbor_objs = retrieval_dataset._data.filter(pc.is_in(retrieval_dataset._data['id'], value_set))
-
-        # build mapping between indices and embeddings
-        neighbor_id_list = neighbor_objs['id']
-        documents = neighbor_objs['text']
-        idx_to_embedding = {idx.as_py(): documents[i] for i, idx in enumerate(neighbor_id_list)}
-
-        neighbor_text = []
-        for cur_neighbor_ids in neighbor_ids:
-            cur_embs = [idx_to_embedding[id].as_py() for id in cur_neighbor_ids]
-            neighbor_text.append(cur_embs)
-
-        neighbor_text_column.extend(neighbor_text)
-        neighbor_ids_column.extend(neighbor_ids)
-
-    print("adding neighbor ids")
-    train_dataset = train_dataset.add_column('neighbor_ids', neighbor_ids_column)
-    print("adding neighbors")
-    train_dataset = train_dataset.add_column('neighbor_text', neighbor_text_column)
-
-    supplemented_dataset_path = f"{ds_path}_supplemented_{args.split}/"
-    train_dataset.save_to_disk(supplemented_dataset_path)
-
-
-if __name__ == "__main__":
-    prep_index()

gpt4all/index/tokenize_texts.py

@@ -1,72 +0,0 @@
-from datasets import load_dataset
-from argparse import ArgumentParser
-from gpt4all.index.embed import Embedder
-
-
-def parse_args():
-    parser = ArgumentParser()
-    # fmt: off
-    parser.add_argument("--tokenized_save_path", type=str, default="tokenized")
-    parser.add_argument("--ds_name", type=str, default="wikipedia")
-    parser.add_argument("--ds_version", type=str, default="20220301.simple")
-    parser.add_argument("--sbert_model", type=str, default="sentence-transformers/all-MiniLM-L6-v2")
-    # fmt: on
-
-    return parser.parse_args()
-
-
-def tokenize_texts(examples, embedder):
-    split_data = {k: [] for k in examples.keys()}
-    split_data["tokenized_chunk"] = []
-    split_data["tokenized_attn_mask"] = []
-
-    keys = [k for k in examples.keys() if k != "text"]
-    for i, text in enumerate(examples["text"]):
-        tokenized_text = embedder.chunk_text(text)
-        # do we want to add sep/cls tokens to beginning and end?
-        decoded_text = embedder.tokenizer.batch_decode(
-            sequences=tokenized_text["input_ids"]
-        )
-
-        num_splits = len(tokenized_text["input_ids"])
-        split_data["id"].extend(
-            [f"{examples['id'][i]}_split_{j}" for j in range(num_splits)]
-        )
-
-        for col in keys:
-            if col != "id":
-                split_data[col].extend(
-                    [examples[col][i]] * len(tokenized_text["input_ids"])
-                )
-
-        split_data["text"].extend(decoded_text)
-        split_data["tokenized_chunk"].extend(tokenized_text["input_ids"])
-        split_data["tokenized_attn_mask"].extend(tokenized_text["attention_mask"])
-
-    return split_data
-
-
-def chunk_dataset(
-    ds_name="wikipedia",
-    version="20220301.simple",
-    sbert_model="sentence-transformers/all-MiniLM-L6-v2",
-    save_path="tokenized",
-):
-    dataset = load_dataset(ds_name, version, split="train")
-    print(len(dataset))
-    embedder = Embedder(sbert_model)
-    dataset = dataset.map(
-        lambda x: tokenize_texts(x, embedder), batched=True, num_proc=64
-    )
-
-    dataset.save_to_disk(save_path)
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    chunked_dataset = chunk_dataset(
-        ds_name=args.ds_name,
-        version=args.ds_version,
-        sbert_model=args.sbert_model,
-        save_path=args.tokenized_save_path,
-    )

gpt4all/inference/combine_synth_data.py

@@ -1,57 +0,0 @@
-import glob
-from argparse import ArgumentParser
-from datasets import Dataset, load_from_disk, concatenate_datasets
-
-
-PROMPT = "Write a question answer pair based on the following context. If the context isn't specific enough, ignore and return 'No question answer pair`. Context : {}\n"
-
-def load_synth_data(data_dir):
-    files = glob.glob(data_dir + "/*")
-
-    ds = concatenate_datasets([load_from_disk(f) for f in files])
-    table = ds.data.table
-
-    filtered = table.filter(table["valid"])
-    ds = Dataset.from_dict(filtered.to_pydict())
-    return ds 
-
-    
-def remove_prompt(examples):
-    outputs = {"text": [], "generated": [], "question": [], "answer": []}
-    for context, generated in zip(examples["text"], examples["generated"]):
-        prompt_w_ctx = PROMPT.format(context)
-        gen_wo_ctx = generated[len(prompt_w_ctx):]
-
-        assert prompt_w_ctx not in gen_wo_ctx
-
-        question = gen_wo_ctx.split("Answer:")[0].replace("Question:", "").strip()
-        answer = gen_wo_ctx.split("Answer:")[1].strip()
-
-        outputs["text"].append(context)
-        outputs["generated"].append(gen_wo_ctx)
-        outputs["question"].append(question)
-        outputs["answer"].append(answer)
-
-    return outputs
-
-        
-
-
-def combine_synth_data(data_dir):
-    ds = load_synth_data(data_dir)
-
-    ds = ds.map(lambda ele: remove_prompt(ele), batched=True, num_proc=64)
-
-    ds.save_to_disk(f"synth_data_combined_{len(ds)/1000:.0f}")
-
-
-
-if __name__ == "__main__":
-    parser = ArgumentParser()
-    parser.add_argument("--dataset_dir", type=str, required=True)
-
-    args = parser.parse_args()
-
-    combine_synth_data(args.dataset_dir)
-
-    

gpt4all/inference/generate_synth_data.py

@@ -1,149 +0,0 @@
-from argparse import ArgumentParser
-from datasets import load_dataset, Dataset
-import torch
-from torch.utils.data import DataLoader, DistributedSampler
-from accelerate.utils import set_seed
-from transformers import AutoTokenizer, AutoModelForCausalLM, DefaultDataCollator 
-from gpt4all.utils.read import read_config
-from gpt4all.utils.distributed_utils import rank0_print, main_process_first
-from tqdm import tqdm
-import pyarrow.compute as pc
-import pyarrow as pa
-import torch.distributed as dist
-
-
-PROMPT = "Write a question answer pair based on the following context. If the context isn't specific enough, ignore and return 'No question answer pair`. Context : {}\n"
-
-
-def prepare_data(config, tokenizer, num_processes, local_rank):
-    dataset = load_dataset(config["dataset_path"], split="train")
-    dataset = dataset.remove_columns("embedding")
-
-    shuffled = dataset.shuffle(seed=config["seed"])
-    indices = shuffled[:config["max_generations"]]["id"]
-
-    table = dataset.data
-    mask = pc.is_in(table["id"], value_set=pa.array(indices, pa.int32()))
-    filtered_table = table.filter(mask)
-
-    # convert from pyarrow to Dataset
-    orig_dataset = Dataset.from_dict(filtered_table.to_pydict())
-
-    dataset = orig_dataset.map(lambda ele: {"prompted_text": [PROMPT.format(context) for context in ele["text"]]},
-                          batched=True,
-                          num_proc=config["num_proc"] if "num_proc" in config else None)
-
-    dataset = dataset.map(lambda ele: {"prompt_len": [len(prompt) for prompt in ele["prompted_text"]]}, batched=True,
-                          num_proc=config["num_proc"] if "num_proc" in config else None)
-
-    dataset = dataset.sort("prompt_len")
-    dataset = dataset.map(lambda ele: tokenizer(ele["prompted_text"], return_tensors="pt", padding="longest", truncation=True,
-                                                max_length=tokenizer.model_max_length - config["max_new_tokens"]), batched=True,
-                                                batch_size=num_processes * config["batch_size"],
-                         )
-
-    columns_to_keep = ["id", "input_ids", "attention_mask"]
-    col_names_to_rm = [col for col in dataset.column_names if col not in columns_to_keep]
-
-    dataset = dataset.remove_columns(col_names_to_rm)
-
-    sampler = DistributedSampler(
-        dataset,
-        shuffle=False,
-        drop_last=True,
-        num_replicas=num_processes,
-        rank=local_rank,
-    )
-
-    dataloader = DataLoader(
-            dataset,
-            batch_size=config["batch_size"],
-            collate_fn=DefaultDataCollator(),
-            sampler=sampler,
-            drop_last=True
-        )
-
-    return dataloader, orig_dataset
-
-def generate_data(config):
-    set_seed(config["seed"])
-
-    rank0_print(f"World size: {dist.get_world_size()}")
-
-    tokenizer = AutoTokenizer.from_pretrained(config["tokenizer_name"])
-    # since we're doing generation, pad left for autoregressive generation
-    tokenizer.padding_side = "left"
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    num_processes = dist.get_world_size()
-    local_rank = dist.get_rank()
-
-    dataloader, dataset = prepare_data(config, tokenizer, num_processes, local_rank)
-
-    dist.barrier()
-    print(dataset[:10]["id"])
-        
-    model = AutoModelForCausalLM.from_pretrained(config["model_name"],
-                                                       revision=config["revision"] if "revision" in config else None,
-                                                       use_cache=True,
-                                                       torch_dtype=torch.bfloat16,)
-    model.to(f"cuda:{local_rank}")
-
-    synth_data = []
-    valid = []
-    ids = []
-    total_valid = 0
-
-    with torch.no_grad():
-        for i, batch in enumerate(tqdm(dataloader, disable=local_rank != 0)):
-            # keep this simple for now, can add temperature and other sampling techniques later
-            generated = model.generate(batch["input_ids"].to(model.device),
-                                       attention_mask=batch["attention_mask"].to(model.device),
-                                       max_new_tokens=config["max_new_tokens"])
-
-            decoded = tokenizer.batch_decode(generated, skip_special_tokens=True)
-            num_valid = ["\nQuestion:" in t and "\nAnswer:" in t for t in decoded]
-            rank0_print(f"Num valid: {sum(num_valid)/ len(num_valid):.2f}")
-            total_valid += sum(num_valid)
-
-            synth_data.extend(decoded)
-            valid.extend(num_valid)
-            ids.extend(batch["id"].tolist())
-
-            if i > 0 and i % config["save_every"] == 0:
-                table = dataset.data.table
-                mask = pc.is_in(table["id"], value_set=pa.array(ids, pa.int32()))
-                filtered_table = table.filter(mask)
-
-                chunk_table = pa.Table.from_pydict({"id": ids, "generated": synth_data, "valid": valid})
-                joined = filtered_table.join(chunk_table, "id")
-                curr_dataset = Dataset.from_dict(joined.to_pydict())
-                curr_dataset.save_to_disk(f'{config["output_path"]}/chunk_{i}_rank_{local_rank}')
-
-    table = dataset.data.table
-    mask = pc.is_in(table["id"], value_set=pa.array(ids, pa.int32()))
-    filtered_table = table.filter(mask)
-
-    chunk_table = pa.Table.from_pydict({"id": ids, "generated": synth_data, "valid": valid})
-    joined = filtered_table.join(chunk_table, "id")
-    full_dataset = Dataset.from_dict(joined.to_pydict())
-    full_dataset.save_to_disk(f'{config["output_path"]}_{config["max_generations"]}_rank_{local_rank}')
-
-    rank0_print(f"Total valid: {total_valid}/{config['max_generations']}")
-
-    
-def main():
-    dist.init_process_group("nccl")
-    parser = ArgumentParser()
-    parser.add_argument("--config", type=str, default="config.yaml")
-
-    args = parser.parse_args()
-    config = read_config(args.config)
-
-    generate_data(config)
-
-
-if __name__ == "__main__":
-    # parse arguments by reading in a config
-    main()

gpt4all/models/__init__.py

@@ -1,18 +0,0 @@
-from .gpt_jr.configuration_gpt_jr import GPTJRConfig
-from .gpt_jr.modeling_gpt_jr import GPTJRForCausalLM
-
-from .pythiaseek import PythiaSeekForCausalLM, PythiaSeekConfig
-from .pythia_retro import PythiaRetroForCausalLM, PythiaRetroConfig
-from .lethe import LetheConfig, LetheForCausalLM
-
-
-__all__ = [
-    "GPTJRConfig",
-    "GPTJRForCausalLM",
-    "PythiaSeekConfig",
-    "PythiaSeekForCausalLM",
-    "PythiaRetroConfig",
-    "PythiaRetroForCausalLM",
-    "LetheConfig",
-    "LetheForCausalLM"
-]

gpt4all/models/gpt_jr/configuration_gpt_jr.py

@@ -1,151 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPT-J model configuration"""
-from collections import OrderedDict
-from typing import Any, List, Mapping, Optional
-
-from transformers import PreTrainedTokenizer, TensorType, is_torch_available
-from transformers.configuration_utils import PretrainedConfig
-from transformers.onnx import OnnxConfigWithPast, PatchingSpec
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "EleutherAI/gpt-j-6B": "https://huggingface.co/EleutherAI/gpt-j-6B/resolve/main/config.json",
-    # See all GPT-J models at https://huggingface.co/models?filter=gpt_j
-}
-
-
-class GPTJRConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`GPTJModel`]. It is used to instantiate a GPT-J
-    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
-    defaults will yield a similar configuration to that of the GPT-J
-    [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B) architecture. Configuration objects inherit from
-    [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`]
-    for more information.
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 50400):
-            Vocabulary size of the GPT-J model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`GPTJModel`].
-        n_positions (`int`, *optional*, defaults to 2048):
-            The maximum sequence length that this model might ever be used with. Typically set this to something large
-            just in case (e.g., 512 or 1024 or 2048).
-        n_embd (`int`, *optional*, defaults to 4096):
-            Dimensionality of the embeddings and hidden states.
-        n_layer (`int`, *optional*, defaults to 28):
-            Number of hidden layers in the Transformer encoder.
-        n_head (`int`, *optional*, defaults to 16):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        rotary_dim (`int`, *optional*, defaults to 64):
-            Number of dimensions in the embedding that Rotary Position Embedding is applied to.
-        n_inner (`int`, *optional*, defaults to None):
-            Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
-        activation_function (`str`, *optional*, defaults to `"gelu_new"`):
-            Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`.
-        resid_pdrop (`float`, *optional*, defaults to 0.1):
-            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
-        embd_pdrop (`int`, *optional*, defaults to 0.1):
-            The dropout ratio for the embeddings.
-        attn_pdrop (`float`, *optional*, defaults to 0.1):
-            The dropout ratio for the attention.
-        layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
-            The epsilon to use in the layer normalization layers.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        scale_attn_weights (`bool`, *optional*, defaults to `True`):
-            Scale attention weights by dividing by sqrt(hidden_size).
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models).
-
-    Example:
-
-    ```python
-    >>> from transformers import GPTJModel, GPTJConfig
-
-    >>> # Initializing a GPT-J 6B configuration
-    >>> configuration = GPTJConfig()
-
-    >>> # Initializing a model from the configuration
-    >>> model = GPTJModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
-    model_type = "gptj"
-    attribute_map = {
-        "max_position_embeddings": "n_positions",
-        "hidden_size": "n_embd",
-        "num_attention_heads": "n_head",
-        "num_hidden_layers": "n_layer",
-    }
-
-    def __init__(
-        self,
-        vocab_size=50400,
-        n_positions=2048,
-        n_embd=4096,
-        n_layer=28,
-        n_head=16,
-        rotary_dim=64,
-        n_inner=None,
-        activation_function="gelu_new",
-        resid_pdrop=0.0,
-        embd_pdrop=0.0,
-        attn_pdrop=0.0,
-        layer_norm_epsilon=1e-5,
-        initializer_range=0.02,
-        scale_attn_weights=True,
-        use_cache=True,
-        bos_token_id=50256,
-        eos_token_id=50256,
-        tie_word_embeddings=False,
-        encoder_dim=4096,
-        total_alpha_steps=0,
-        initial_alpha=1,
-        final_alpha=.5,
-        **kwargs
-    ):
-        self.vocab_size = vocab_size
-        self.n_positions = n_positions
-        self.n_embd = n_embd
-        self.n_layer = n_layer
-        self.n_head = n_head
-        self.n_inner = n_inner
-        self.rotary_dim = rotary_dim
-        self.activation_function = activation_function
-        self.resid_pdrop = resid_pdrop
-        self.embd_pdrop = embd_pdrop
-        self.attn_pdrop = attn_pdrop
-        self.layer_norm_epsilon = layer_norm_epsilon
-        self.initializer_range = initializer_range
-        self.scale_attn_weights = scale_attn_weights
-        self.use_cache = use_cache
-
-        self.bos_token_id = bos_token_id
-        self.eos_token_id = eos_token_id
-
-        self.encoder_dim = encoder_dim
-
-        self.total_alpha_steps = total_alpha_steps
-        self.initial_alpha = initial_alpha
-        self.final_alpha = final_alpha
-
-        super().__init__(
-            bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
-        )

gpt4all/models/gpt_jr/modeling_gpt_jr.py

@@ -1,953 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch GPT-J model."""
-
-from typing import Optional, Tuple, Union
-
-import math
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from transformers import AutoModel
-from transformers.activations import ACT2FN
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import logging
-from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
-from gpt4all.models.gpt_jr.configuration_gpt_jr import GPTJRConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-GPTJR_PRETRAINED_MODEL_ARCHIVE_LIST = [
-    "EleutherAI/gpt-j-6B",
-    # See all GPT-J models at https://huggingface.co/models?filter=gptj
-]
-
-
-def fixed_pos_embedding(x, seq_dim=1, seq_len=None):
-    dim = x.shape[-1]
-    if seq_len is None:
-        seq_len = x.shape[seq_dim]
-    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))
-    sinusoid_inp = (
-        torch.einsum("i , j -> i j", torch.arange(seq_len, dtype=torch.float), inv_freq).to(x.device).float()
-    )
-    return torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)
-
-
-def rotate_every_two(x):
-    x1 = x[:, :, :, ::2]
-    x2 = x[:, :, :, 1::2]
-    x = torch.stack((-x2, x1), dim=-1)
-    return x.flatten(-2)  # in einsum notation: rearrange(x, '... d j -> ... (d j)')
-
-
-def duplicate_interleave(m):
-    """
-    A simple version of `torch.repeat_interleave` for duplicating a matrix while interleaving the copy.
-    """
-    dim0 = m.shape[0]
-    m = m.view(-1, 1)  # flatten the matrix
-    m = m.repeat(1, 2)  # repeat all elements into the 2nd dimension
-    m = m.view(dim0, -1)  # reshape into a matrix, interleaving the copy
-    return m
-
-
-def apply_rotary_pos_emb(x, sincos, offset=0):
-    sin, cos = map(lambda t: duplicate_interleave(t)[None, offset : x.shape[1] + offset, None, :], sincos)
-    # einsum notation for lambda t: repeat(t[offset:x.shape[1]+offset,:], "n d -> () n () (d j)", j=2)
-    return (x * cos) + (rotate_every_two(x) * sin)
-
-    
-
-
-class GPTJRAttention(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-
-        max_positions = config.max_position_embeddings
-        self.register_buffer(
-            "bias",
-            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
-                1, 1, max_positions, max_positions
-            ),
-        )
-        self.register_buffer("masked_bias", torch.tensor(-1e9))
-
-        self.attn_dropout = nn.Dropout(config.attn_pdrop)
-        self.resid_dropout = nn.Dropout(config.resid_pdrop)
-
-        self.embed_dim = config.hidden_size
-        self.num_attention_heads = config.num_attention_heads
-        self.head_dim = self.embed_dim // self.num_attention_heads
-        if self.head_dim * self.num_attention_heads != self.embed_dim:
-            raise ValueError(
-                f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
-                f" `num_attention_heads`: {self.num_attention_heads})."
-            )
-        self.scale_attn = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype())
-
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.rotary_dim = None
-        if config.rotary_dim is not None:
-            self.rotary_dim = config.rotary_dim
-
-    def _split_heads(self, tensor, num_attention_heads, attn_head_size, rotary):
-        """
-        Splits hidden dim into attn_head_size and num_attention_heads
-        """
-        new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
-        tensor = tensor.view(new_shape)
-        if rotary:
-            return tensor
-        if len(tensor.shape) == 5:
-            return tensor.permute(0, 1, 3, 2, 4)  # (batch, blocks, head, block_length, head_features)
-        elif len(tensor.shape) == 4:
-            return tensor.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
-        else:
-            raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
-
-    def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
-        """
-        Merges attn_head_size dim and num_attn_heads dim into hidden dim
-        """
-        if len(tensor.shape) == 5:
-            tensor = tensor.permute(0, 1, 3, 2, 4).contiguous()
-        elif len(tensor.shape) == 4:
-            tensor = tensor.permute(0, 2, 1, 3).contiguous()
-        else:
-            raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
-        new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
-        return tensor.view(new_shape)
-
-    def _attn(
-        self,
-        query,
-        key,
-        value,
-        attention_mask=None,
-        head_mask=None,
-    ):
-
-        # compute causal mask from causal mask buffer
-        query_length, key_length = query.size(-2), key.size(-2)
-        causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].to(torch.bool)
-
-        attn_weights = torch.matmul(query, key.transpose(-1, -2))
-
-        mask_value = torch.finfo(attn_weights.dtype).min
-        # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
-        # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
-        mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
-        attn_weights = torch.where(causal_mask, attn_weights, mask_value)
-
-        attn_weights = attn_weights / self.scale_attn
-
-        if attention_mask is not None:
-            # Apply the attention mask
-            attn_weights = attn_weights + attention_mask
-
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-        attn_weights = attn_weights.to(value.dtype)
-        attn_weights = self.attn_dropout(attn_weights)
-
-        # Mask heads if we want to
-        if head_mask is not None:
-            attn_weights = attn_weights * head_mask
-
-        attn_output = torch.matmul(attn_weights, value)
-
-        return attn_output, attn_weights
-
-    def forward(
-        self,
-        hidden_states: Optional[torch.FloatTensor],
-        encoder_hidden_states: Optional[torch.FloatTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        layer_past: Optional[Tuple[torch.Tensor]] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = False,
-        output_attentions: Optional[bool] = False,
-    ) -> Union[
-        Tuple[torch.Tensor, Tuple[torch.Tensor]],
-        Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
-    ]:
-
-        query = self.q_proj(hidden_states)
-        # if we are doing cross attention
-        if encoder_hidden_states is not None:
-            key = self.k_proj(encoder_hidden_states)
-            value = self.v_proj(encoder_hidden_states)
-        else:
-            key = self.k_proj(hidden_states)
-            value = self.v_proj(hidden_states)
-
-        query = self._split_heads(query, self.num_attention_heads, self.head_dim, True)
-        key = self._split_heads(key, self.num_attention_heads, self.head_dim, True)
-        value = self._split_heads(value, self.num_attention_heads, self.head_dim, False)
-
-        seq_len = key.shape[1]
-        offset = 0
-
-        if layer_past is not None:
-            offset = layer_past[0].shape[-2]
-            seq_len += offset
-
-        if self.rotary_dim is not None:
-            k_rot = key[:, :, :, : self.rotary_dim]
-            k_pass = key[:, :, :, self.rotary_dim :]
-
-            q_rot = query[:, :, :, : self.rotary_dim]
-            q_pass = query[:, :, :, self.rotary_dim :]
-
-            sincos = fixed_pos_embedding(k_rot, 1, seq_len=seq_len)
-            k_rot = apply_rotary_pos_emb(k_rot, sincos, offset=offset)
-            q_rot = apply_rotary_pos_emb(q_rot, sincos, offset=offset)
-
-            key = torch.cat([k_rot, k_pass], dim=-1)
-            query = torch.cat([q_rot, q_pass], dim=-1)
-        else:
-            sincos = fixed_pos_embedding(key, 1, seq_len=seq_len)
-            key = apply_rotary_pos_emb(key, sincos, offset=offset)
-            query = apply_rotary_pos_emb(query, sincos, offset=offset)
-
-        key = key.permute(0, 2, 1, 3)
-        query = query.permute(0, 2, 1, 3)
-
-        if layer_past is not None:
-            past_key = layer_past[0]
-            past_value = layer_past[1]
-            key = torch.cat((past_key, key), dim=-2)
-            value = torch.cat((past_value, value), dim=-2)
-
-        if use_cache is True:
-            present = (key, value)
-        else:
-            present = None
-
-        # compute self-attention: V x Softmax(QK^T)
-        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
-        attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
-        attn_output = self.out_proj(attn_output)
-        attn_output = self.resid_dropout(attn_output)
-
-        outputs = (attn_output, present)
-        if output_attentions:
-            outputs += (attn_weights,)
-
-        return outputs  # a, present, (attentions)
-
-        
-class GPTJRCrossAttention(GPTJRAttention):
-    def __init__(self, config):
-        super().__init__(config)
-
-        max_positions = config.max_position_embeddings
-        self.register_buffer(
-            "bias",
-            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
-                1, 1, max_positions, max_positions
-            ),
-        )
-        self.register_buffer("masked_bias", torch.tensor(-1e9))
-
-        self.attn_dropout = nn.Dropout(config.attn_pdrop)
-        self.resid_dropout = nn.Dropout(config.resid_pdrop)
-
-        self.embed_dim = config.hidden_size
-        self.num_attention_heads = config.num_attention_heads
-        self.head_dim = self.embed_dim // self.num_attention_heads
-        if self.head_dim * self.num_attention_heads != self.embed_dim:
-            raise ValueError(
-                f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
-                f" `num_attention_heads`: {self.num_attention_heads})."
-            )
-        self.scale_attn = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype())
-
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-
-        
-    def _split_knn_attn_heads(self, tensor, num_attention_heads, attn_head_size):
-        new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
-        tensor = tensor.view(new_shape)
-
-        return tensor.permute(0, 1, 3, 2)
-
-        
-    def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
-        """
-        Merges attn_head_size dim and num_attn_heads dim into hidden dim
-        """
-        # tensor -> (bs, seq_len, num_attention_heads, head_dim)
-        tensor = tensor.permute(0, 2, 1, 3).contiguous()
-        new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
-        return tensor.view(new_shape)
-
-    def _attn(
-        self,
-        query,
-        key,
-        value,
-        attention_mask=None,
-        head_mask=None,
-    ):
-
-        # query -> (bs, num_attention_heads, seq_len, head_dim)
-        # key -> (bs, num_attention_heads, head_dim, neighbors)
-        # attn_weights -> (bs, num_attention_heads, seq_len, neighbors)
-        attn_weights = torch.matmul(query, key)
-
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-        attn_weights = attn_weights.to(value.dtype)
-        attn_weights = self.attn_dropout(attn_weights)
-
-        # Mask heads if we want to
-        if head_mask is not None:
-            attn_weights = attn_weights * head_mask
-
-        # value -> (bs, num_attention_heads, seq_len, head_dim)
-        # attn_weights -> (bs, num_attention_heads, seq_len, neighbors)
-        # attn_output -> (bs, num_attention_heads, seq_len, head_dim)
-        attn_output = torch.matmul(attn_weights, value)
-
-        return attn_output, attn_weights
-
-            
-    def forward(
-        self,
-        hidden_states: Optional[torch.FloatTensor],
-        encoder_hidden_states: Optional[torch.FloatTensor],
-        attention_mask: Optional[torch.FloatTensor] = None,
-        layer_past: Optional[Tuple[torch.Tensor]] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = False,
-        output_attentions: Optional[bool] = False,
-    ) -> Union[
-        Tuple[torch.Tensor, Tuple[torch.Tensor]],
-        Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
-    ]:
-
-        query = self.q_proj(hidden_states)
-        # if we are doing cross attention
-        key = self.k_proj(encoder_hidden_states)
-        value = self.v_proj(encoder_hidden_states)
-        # (bs, seq_len, dim) -> (bs, num_attention_heads, seq_len, head_dim)
-        query = self._split_heads(query, self.num_attention_heads, self.head_dim, False)
-        # (bs, dim) -> (bs, num_attention_heads,  head_dim)
-        key = self._split_knn_attn_heads(key, self.num_attention_heads, self.head_dim)
-        value = self._split_knn_attn_heads(value, self.num_attention_heads, self.head_dim)
-
-
-        value = value.permute(0, 3, 1, 2)
-        key = key.permute(0, 3, 2, 1)
-
-        if layer_past is not None:
-            past_key = layer_past[0]
-            past_value = layer_past[1]
-            key = torch.cat((past_key, key), dim=-2)
-            value = torch.cat((past_value, value), dim=-2)
-
-        if use_cache is True:
-            present = (key, value)
-        else:
-            present = None
-
-        # compute self-attention: V x Softmax(QK^T)
-        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
-        attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
-        attn_output = self.out_proj(attn_output)
-        attn_output = self.resid_dropout(attn_output)
-
-        outputs = (attn_output, present)
-        if output_attentions:
-            outputs += (attn_weights,)
-
-        return outputs  # a, present, (attentions)
-        
-
-
-class GPTJRMLP(nn.Module):
-    def __init__(self, intermediate_size, config):  # in MLP: intermediate_size= 4 * embed_dim
-        super().__init__()
-        embed_dim = config.n_embd
-
-        self.fc_in = nn.Linear(embed_dim, intermediate_size)
-        self.fc_out = nn.Linear(intermediate_size, embed_dim)
-
-        self.act = ACT2FN[config.activation_function]
-        self.dropout = nn.Dropout(config.resid_pdrop)
-
-    def forward(self, hidden_states: Optional[torch.FloatTensor]) -> torch.FloatTensor:
-        hidden_states = self.fc_in(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.fc_out(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        return hidden_states
-
-
-class GPTJRBlock(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
-        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-        self.attn = GPTJRAttention(config)
-        self.mlp = GPTJRMLP(inner_dim, config)
-
-        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-        self.cross_attn = GPTJRCrossAttention(config)
-        self.cross_attn_mlp = GPTJRMLP(inner_dim, config)
-        
-        self.total_alpha_steps = config.total_alpha_steps
-        self.initial_alpha = config.initial_alpha
-        self.final_alpha = config.final_alpha
-
-    def forward(
-        self,
-        hidden_states: Optional[torch.FloatTensor],
-        encoder_hidden_states: Optional[torch.FloatTensor],
-        layer_past: Optional[Tuple[torch.Tensor]] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = False,
-        output_attentions: Optional[bool] = False,
-        step: Optional[int] = None,
-    ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
-        # shape (bs, seq_len, hidden_dim)
-        residual = hidden_states
-        hidden_states = self.ln_1(hidden_states)
-        attn_outputs = self.attn(
-            hidden_states,
-            layer_past=layer_past,
-            attention_mask=attention_mask,
-            head_mask=head_mask,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-        )
-        attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
-        outputs = attn_outputs[1:]
-
-        feed_forward_hidden_states = self.mlp(hidden_states)
-        self_attention_residual = attn_output + feed_forward_hidden_states + residual
-
-        # encoder_hidden_states -> (bs, knn, encoder_dim)
-        if encoder_hidden_states.dtype != hidden_states.dtype:
-            encoder_hidden_states = encoder_hidden_states.to(hidden_states.dtype)
-        encoder_normed = self.ln_2(encoder_hidden_states)
-
-        # cross_attn_outputs -> (bs, seq_len, dim)
-        cross_attn_output = self.cross_attn(
-            hidden_states,
-            encoder_hidden_states=encoder_normed,
-            attention_mask=attention_mask,
-            head_mask=head_mask,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-        )
-
-        # gpt-j has parallel ff + attn, can do ff on encoder_normed too I guess?
-        cross_attn_ff = self.cross_attn_mlp(
-            cross_attn_output[0] 
-        )
-
-        if step is not None:
-            alpha = self._update_alpha(step)
-        else:
-            alpha = 0.5 
-
-        hidden_states = (1 - alpha) * cross_attn_ff +  alpha * self_attention_residual  
-
-        if use_cache:
-            outputs = (hidden_states,) + outputs
-        else:
-            outputs = (hidden_states,) + outputs[1:]
-
-        return outputs  # hidden_states, present, (attentions)
-        
-    def _update_alpha(self, current_step):
-        """
-        Computes the learning rate for the current step using a cosine decay schedule.
-
-        Args:
-            initial_lr (float): The initial learning rate.
-            final_lr (float): The final learning rate.
-            total_steps (int): The total number of steps in the schedule.
-            current_step (int): The current step.
-
-        Returns:
-            float: The learning rate for the current step.
-        """
-        initial_alpha = 1
-        final_alpha = .5
-        if current_step >= self.total_alpha_steps:
-            return final_alpha
-
-        # Compute the cosine decay factor
-        cosine_decay = 0.5 * (1 + math.cos(math.pi * current_step / self.total_alpha_steps))
-
-        # Compute the current learning rate
-        alpha = final_alpha + (initial_alpha - final_alpha) * cosine_decay
-
-        return alpha
-
-
-class GPTJRPreTrainedModel(PreTrainedModel):
-    """
-    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
-    models.
-    """
-
-    config_class = GPTJRConfig
-    base_model_prefix = "transformer"
-    is_parallelizable = True
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["GPTJRBlock"]
-
-    def __init__(self, *inputs, **kwargs):
-        super().__init__(*inputs, **kwargs)
-
-    def _init_weights(self, module):
-        """Initialize the weights."""
-        if isinstance(module, (nn.Linear,)):
-            # Slightly different from Mesh Transformer JAX which uses truncated_normal for initialization
-            # cf https://github.com/pytorch/pytorch/pull/5617
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, GPTJRModel):
-            module.gradient_checkpointing = value
-
-
-class GPTJRModel(GPTJRPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.embed_dim = config.n_embd
-        self.vocab_size = config.vocab_size
-        self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
-        self.drop = nn.Dropout(config.embd_pdrop)
-        self.h = nn.ModuleList([GPTJRBlock(config) for _ in range(config.n_layer)])
-        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
-
-        # Model parallel
-        self.model_parallel = False
-        self.device_map = None
-        self.gradient_checkpointing = False
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def parallelize(self, device_map=None):
-        # Check validity of device_map
-        self.device_map = (
-            get_device_map(len(self.h), range(torch.cuda.device_count())) if device_map is None else device_map
-        )
-        assert_device_map(self.device_map, len(self.h))
-        self.model_parallel = True
-        self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys()))
-        self.last_device = "cuda:" + str(max(self.device_map.keys()))
-        self.wte = self.wte.to(self.first_device)
-        # Load onto devices
-        for k, v in self.device_map.items():
-            for block in v:
-                cuda_device = "cuda:" + str(k)
-                self.h[block] = self.h[block].to(cuda_device)
-        # ln_f to last
-        self.ln_f = self.ln_f.to(self.last_device)
-
-    def deparallelize(self):
-        self.model_parallel = False
-        self.device_map = None
-        self.first_device = "cpu"
-        self.last_device = "cpu"
-        self.wte = self.wte.to("cpu")
-        for index in range(len(self.h)):
-            self.h[index] = self.h[index].to("cpu")
-        self.ln_f = self.ln_f.to("cpu")
-        torch.cuda.empty_cache()
-
-    def get_input_embeddings(self):
-        return self.wte
-
-    def set_input_embeddings(self, new_embeddings):
-        self.wte = new_embeddings
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        encoder_hidden_states: Optional[torch.FloatTensor] = None,
-        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        step: Optional[int] = None,
-    ) -> Union[Tuple, BaseModelOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
-        elif input_ids is not None:
-            input_shape = input_ids.size()
-            input_ids = input_ids.view(-1, input_shape[-1])
-            batch_size = input_ids.shape[0]
-        elif inputs_embeds is not None:
-            input_shape = inputs_embeds.size()[:-1]
-            batch_size = inputs_embeds.shape[0]
-        else:
-            raise ValueError("You have to specify either input_ids or inputs_embeds")
-
-        device = input_ids.device if input_ids is not None else inputs_embeds.device
-
-        if token_type_ids is not None:
-            token_type_ids = token_type_ids.view(-1, input_shape[-1])
-
-        if position_ids is not None:
-            position_ids = position_ids.view(-1, input_shape[-1])
-
-        if past_key_values is None:
-            past_length = 0
-            past_key_values = tuple([None] * len(self.h))
-        else:
-            past_length = past_key_values[0][0].size(-2)
-
-        if position_ids is None:
-            position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
-            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
-
-        # Attention mask.
-        if attention_mask is not None:
-            if batch_size <= 0:
-                raise ValueError("batch_size has to be defined and > 0")
-            attention_mask = attention_mask.view(batch_size, -1)
-            # We create a 3D attention mask from a 2D tensor mask.
-            # Sizes are [batch_size, 1, 1, to_seq_length]
-            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
-            # this attention mask is more simple than the triangular masking of causal attention
-            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
-            attention_mask = attention_mask[:, None, None, :]
-
-            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
-            # masked positions, this operation will create a tensor which is 0.0 for
-            # positions we want to attend and the dtype's smallest value for masked positions.
-            # Since we are adding it to the raw scores before the softmax, this is
-            # effectively the same as removing these entirely.
-            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
-            attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
-        # Prepare head mask if needed
-        # 1.0 in head_mask indicate we keep the head
-        # attention_probs has shape bsz x num_attention_heads x N x N
-        # head_mask has shape n_layer x batch x num_attention_heads x N x N
-        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
-        if inputs_embeds is None:
-            inputs_embeds = self.wte(input_ids)
-
-        hidden_states = inputs_embeds
-
-        if token_type_ids is not None:
-            token_type_embeds = self.wte(token_type_ids)
-            hidden_states = hidden_states + token_type_embeds
-
-        hidden_states = self.drop(hidden_states)
-
-        output_shape = input_shape + (hidden_states.size(-1),)
-
-        presents = () if use_cache else None
-        all_self_attentions = () if output_attentions else None
-        all_hidden_states = () if output_hidden_states else None
-        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
-
-            # Model parallel
-            if self.model_parallel:
-                torch.cuda.set_device(hidden_states.device)
-                # Ensure layer_past is on same device as hidden_states (might not be correct)
-                if layer_past is not None:
-                    layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
-                # Ensure that attention_mask is always on the same device as hidden_states
-                if attention_mask is not None:
-                    attention_mask = attention_mask.to(hidden_states.device)
-                if isinstance(head_mask, torch.Tensor):
-                    head_mask = head_mask.to(hidden_states.device)
-            if output_hidden_states:
-                all_hidden_states = all_hidden_states + (hidden_states,)
-
-            if self.gradient_checkpointing and self.training:
-
-                if use_cache:
-                    logger.warning(
-                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                    )
-                    use_cache = False
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs, use_cache, output_attentions, step)
-
-                    return custom_forward
-
-                outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    encoder_hidden_states,
-                    None,
-                    attention_mask,
-                    head_mask[i],
-                )
-            else:
-                outputs = block(
-                    hidden_states,
-                    encoder_hidden_states,
-                    layer_past=layer_past,
-                    attention_mask=attention_mask,
-                    head_mask=head_mask[i],
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                    step=step
-                )
-
-            hidden_states = outputs[0]
-            if use_cache is True:
-                presents = presents + (outputs[1],)
-
-            if output_attentions:
-                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
-
-            # Model Parallel: If it's the last layer for that device, put things on the next device
-            if self.model_parallel:
-                for k, v in self.device_map.items():
-                    if i == v[-1] and "cuda:" + str(k) != self.last_device:
-                        hidden_states = hidden_states.to("cuda:" + str(k + 1))
-
-        hidden_states = self.ln_f(hidden_states)
-
-        hidden_states = hidden_states.view(output_shape)
-        # Add last hidden state
-        if output_hidden_states:
-            all_hidden_states = all_hidden_states + (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
-
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=presents,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attentions,
-        )
-
-
-class GPTJRForCausalLM(GPTJRPreTrainedModel):
-    _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.masked_bias", r"h\.\d+\.attn\.bias"]
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.transformer = GPTJRModel(config)
-        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
-
-        self.hidden_size = config.hidden_size
-        self.encoder_dim = config.encoder_dim
-
-        if self.hidden_size != self.encoder_dim:
-            self.enc_dec_proj = nn.Sequential(nn.Linear(config.encoder_dim, config.n_embd * 4),
-                                              nn.Linear(config.n_embd * 4, config.n_embd))
-
-        # Model parallel
-        self.model_parallel = False
-        self.device_map = None
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def parallelize(self, device_map=None):
-        self.device_map = (
-            get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
-            if device_map is None
-            else device_map
-        )
-        assert_device_map(self.device_map, len(self.transformer.h))
-        self.transformer.parallelize(self.device_map)
-        self.lm_head = self.lm_head.to(self.transformer.first_device)
-        self.model_parallel = True
-
-    def deparallelize(self):
-        self.transformer.deparallelize()
-        self.transformer = self.transformer.to("cpu")
-        self.lm_head = self.lm_head.to("cpu")
-        self.model_parallel = False
-        torch.cuda.empty_cache()
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
-        token_type_ids = kwargs.get("token_type_ids", None)
-        # only last token for inputs_ids if past is defined in kwargs
-        if past_key_values:
-            input_ids = input_ids[:, -1].unsqueeze(-1)
-            if token_type_ids is not None:
-                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
-
-        attention_mask = kwargs.get("attention_mask", None)
-        position_ids = kwargs.get("position_ids", None)
-
-        if attention_mask is not None and position_ids is None:
-            # create position_ids on the fly for batch generation
-            position_ids = attention_mask.long().cumsum(-1) - 1
-            position_ids.masked_fill_(attention_mask == 0, 1)
-            if past_key_values:
-                position_ids = position_ids[:, -1].unsqueeze(-1)
-        else:
-            position_ids = None
-        return {
-            "input_ids": input_ids,
-            "past_key_values": past_key_values,
-            "use_cache": kwargs.get("use_cache"),
-            "position_ids": position_ids,
-            "attention_mask": attention_mask,
-            "token_type_ids": token_type_ids,
-        }
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        encoder_hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        step: Optional[int] = None
-    ) -> Union[Tuple, CausalLMOutputWithPast]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
-            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
-            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
-        """
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if self.hidden_size != self.encoder_dim:
-            encoder_hidden_states = encoder_hidden_states.to(self.enc_dec_proj[0].weight.dtype)
-            encoder_hidden_states = self.enc_dec_proj(encoder_hidden_states)
-
-        transformer_outputs = self.transformer(
-            input_ids,
-            encoder_hidden_states=encoder_hidden_states,
-            past_key_values=past_key_values,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            step=step,
-        )
-        hidden_states = transformer_outputs[0]
-
-        # Set device for model parallelism
-        if self.model_parallel:
-            torch.cuda.set_device(self.transformer.first_device)
-            hidden_states = hidden_states.to(self.lm_head.weight.device)
-
-        # make sure sampling in fp16 works correctly and
-        # compute loss in fp32 to match with mesh-tf version
-        # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
-        # TODO: do we need to do conversion to fp32 if training in bf16?
-        lm_logits = self.lm_head(hidden_states).to(torch.float32)
-
-        loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = lm_logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
-            loss = loss.to(hidden_states.dtype)
-
-        if not return_dict:
-            output = (lm_logits,) + transformer_outputs[1:]
-            return ((loss,) + output) if loss is not None else output
-
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=lm_logits,
-            past_key_values=transformer_outputs.past_key_values,
-            hidden_states=transformer_outputs.hidden_states,
-            attentions=transformer_outputs.attentions,
-        )
-
-    @staticmethod
-    def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
-        """
-        This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
-        [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
-        beam_idx at every generation step.
-        """
-        return tuple(
-            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
-            for layer_past in past
-        )
-

gpt4all/models/gpt_jr/test_gpt_jr.py

@@ -1,50 +0,0 @@
-import torch
-from gpt4all.models import GPTJRForCausalLM, GPTJRConfig
-from transformers import AutoTokenizer, AutoModel
-
-# seed torch 
-
-torch.manual_seed(0)
-
-config = GPTJRConfig(encoder_dim=384, n_layer=4)
-print("loaded config")
-
-print("loading model")
-model = GPTJRForCausalLM(config)
-print("loaded model")
-
-
-tokenizer =  AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6b")
-tokenizer.pad_token = tokenizer.eos_token
-
-encoder_tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-encoder = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-
-
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-
-text = "The quick brown fox jumps over the lazy dog."
-print("Encoded knn")
-tokenized = encoder_tokenizer(text, return_tensors="pt")
-
-# bs, seq_len, dim
-encodings = mean_pooling(encoder(**tokenized), tokenized["attention_mask"])
-
-# make 2 neighbors
-# (bs, knn, encoding_dim)
-encoder_hidden_states = torch.stack([encodings, encodings]).squeeze().unsqueeze(0)
-
-inputs = "What did the fox do?"
-
-print("Encoded inputs")
-tokenized_input = tokenizer([inputs], padding="max_length", truncation=True, return_tensors="pt")
-
-print("Running model")
-outputs = model(**tokenized_input, encoder_hidden_states=encoder_hidden_states)
-
-print(outputs)
-print(outputs[0].shape)
-

gpt4all/models/lethe/__init__.py

@@ -1,2 +0,0 @@
-from .configuration_lethe import LetheConfig
-from .modeling_lethe import LetheForCausalLM

gpt4all/models/lethe/configuration_lethe.py

@@ -1,138 +0,0 @@
-# coding=utf-8
-# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPTNeoX model configuration"""
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "EleutherAI/gpt-neox-20b": "https://huggingface.co/EleutherAI/gpt-neox-20b/resolve/main/config.json",
-    # See all GPTNeoX models at https://huggingface.co/models?filter=gpt_neox
-}
-
-
-class LetheConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`GPTNeoXModel`]. It is used to instantiate an
-    GPTNeoX model according to the specified arguments, defining the model architecture. Instantiating a configuration
-    with the defaults will yield a similar configuration to that of the GPTNeoX
-    [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 50432):
-            Vocabulary size of the GPTNeoX model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`GPTNeoXModel`].
-        hidden_size (`int`, *optional*, defaults to 6144):
-            Dimension of the encoder layers and the pooler layer.
-        num_hidden_layers (`int`, *optional*, defaults to 44):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 64):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        intermediate_size (`int`, *optional*, defaults to 24576):
-            Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
-        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
-            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
-            `"relu"`, `"selu"` and `"gelu_new"` are supported.
-        rotary_pct (`float`, *optional*, defaults to 0.25):
-            percentage of hidden dimensions to allocate to rotary embeddings
-        rotary_emb_base (`int`, *optional*, defaults to 10000)
-            base for computing rotary embeddings frequency
-        classifier_dropout (`float`, *optional*, defaults to 0.1):
-            Argument used when doing token classification, used in the model [`GPTNeoXForTokenClassification`].
-
-            The dropout ratio for the hidden layer.
-        max_position_embeddings (`int`, *optional*, defaults to 2048):
-            The maximum sequence length that this model might ever be used with. Typically set this to something large
-            just in case (e.g., 512 or 1024 or 2048).
-        initializer_range (`float`, *optional*, defaults to 1e-5):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
-            The epsilon used by the layer normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models). Only
-            relevant if `config.is_decoder=True`.
-        use_parallel_residual (`bool`, *optional*, defaults to `True`):
-            Whether to use a "parallel" formulation in each Transformer layer, which can provide a slight training
-            speedup at large scales (e.g. 20B).
-        Example:
-
-    ```python
-    >>> from transformers import GPTNeoXConfig, GPTNeoXModel
-
-    >>> # Initializing a GPTNeoX gpt-neox-20b style configuration
-    >>> configuration = GPTNeoXConfig()
-
-    >>> # Initializing a model (with random weights) from the gpt-neox-20b style configuration
-    >>> model = GPTNeoXModel(configuration)  # doctest: +SKIP
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config  # doctest: +SKIP
-    ```"""
-    model_type = "gpt_neox"
-
-    def __init__(
-        self,
-        vocab_size=50432,
-        hidden_size=6144,
-        num_hidden_layers=44,
-        num_attention_heads=64,
-        intermediate_size=24576,
-        hidden_act="gelu",
-        rotary_pct=0.25,
-        rotary_emb_base=10000,
-        classifier_dropout=0.1,
-        max_position_embeddings=2048,
-        initializer_range=0.02,
-        layer_norm_eps=1e-5,
-        use_cache=True,
-        bos_token_id=0,
-        eos_token_id=2,
-        tie_word_embeddings=False,
-        use_parallel_residual=True,
-        memory_attn_layer=9,
-        num_neighbors_to_retrieve=32,
-        num_neighbors_stored=128,
-        attn_scale_init=20.0,
-        **kwargs,
-    ):
-        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.intermediate_size = intermediate_size
-        self.hidden_act = hidden_act
-        self.rotary_pct = rotary_pct
-        self.rotary_emb_base = rotary_emb_base
-        self.classifier_dropout = classifier_dropout
-        self.initializer_range = initializer_range
-        self.layer_norm_eps = layer_norm_eps
-        self.use_cache = use_cache
-        self.tie_word_embeddings = tie_word_embeddings
-        self.use_parallel_residual = use_parallel_residual
-
-        # index of cross attention layer to add
-        self.memory_attn_layer = memory_attn_layer 
-        self.num_neighbors_to_retrieve = num_neighbors_to_retrieve
-        self.num_neighbors_stored = num_neighbors_stored
-        self.attn_scale_init = attn_scale_init

gpt4all/models/lethe/test_index.py

@@ -1,21 +0,0 @@
-import time
-import numpy as np
-from gpt4all.models.lethe.modeling_lethe import MemoryIndex
-
-
-index = MemoryIndex(256,
-                    575000,
-                    8 
-)
-
-keys = np.random.randn(32, 8, 1024, 256)
-values = np.random.randn(32, 8, 1024, 256)
-start = time.time()
-index.add(keys, values)
-print(f"index.add time: {time.time() - start}")
-
-print(index.key_indices[0].index.ntotal)
-queries = np.random.randn(32, 8, 1024, 256)
-start = time.time()
-index.knn_query(queries, k=32)
-print(f"index.knn_query time: {time.time() - start}")

gpt4all/models/lethe/test_lethe.py

@@ -1,84 +0,0 @@
-import torch
-from gpt4all.models import LetheForCausalLM, LetheConfig
-from gpt4all.models.lethe.modeling_lethe import MemoryIndex
-from transformers import AutoTokenizer, AutoModel
-from datasets import load_from_disk
-from tqdm import tqdm
-
-# seed torch 
-
-torch.manual_seed(0)
-
-config = LetheConfig(num_hidden_layers=15, 
-                     hidden_size=2048, 
-                     intermediate_size=8192, 
-                     num_attention_heads=8,
-                     memory_attn_layer=12,
-                     num_neighbors_stored=6_000_000,
-                     num_neighbors_to_retrieve=32,
-)
-print("loaded config")
-
-print("loading model")
-dimension = config.max_position_embeddings * config.hidden_size
-head_size = config.hidden_size // config.num_attention_heads
-index = MemoryIndex(head_size,
-                    5_000_000,
-                    # 2 since multi-query attention and storing one each for key and value
-                    config.num_attention_heads,
-)
-model = LetheForCausalLM(config, index)
-model.to("cuda:0")
-print("loaded model")
-
-
-tokenizer =  AutoTokenizer.from_pretrained("EleutherAI/pythia-1b")
-tokenizer.pad_token = tokenizer.eos_token
-tokenizer.model_max_length = 2048
-
-dataset = load_from_disk("/home/paperspace/gpt4all/gpt4all/index/squad_supplemented_train")
-
-item = dataset[0]
-
-question = item["question"]
-answer = item["answers"]["text"][0]
-print(f"question: {question}")
-print(f"answer: {answer}")
-
-contexts = item["neighbor_text"]
-
-contexts_encoded = tokenizer(contexts, padding="max_length", truncation=True, return_tensors="pt")
-tokenized_input = tokenizer(question + "\n" + answer, return_tensors="pt", padding="max_length", truncation=True)
-
-inputs = torch.concatenate([contexts_encoded["input_ids"], tokenized_input["input_ids"]], axis=0)
-token_type_ids = torch.tensor([[0] * len(contexts_encoded["input_ids"]) + [1] * len(tokenized_input["input_ids"])]).squeeze()
-
-question_len = len(tokenizer(question + "\n", return_tensors="pt")["input_ids"][0])
-
-labels = inputs.clone()
-labels[:-1] = -100
-labels[-1, :question_len] = -100
-
-
-memory_mask = token_type_ids == 0
-# should be shape (num_memories, sequence_length)
-memories = inputs[memory_mask]
-
-print("Running model")
-model.eval()
-with torch.no_grad():
-    for chunk_start in tqdm(range(0, memories.shape[0], 32)):
-        chunk_end = min(memories.shape[0], chunk_start + 32)
-        mem_chunk = memories[chunk_start:chunk_end].to(model.device)
-        model(input_ids=mem_chunk, labels=None,)
-
-model.train()
-
-qa_inputs = inputs[~memory_mask]
-qa_labels = labels[~memory_mask]
-outputs = model(input_ids=qa_inputs.to(model.device), labels=qa_labels.to(model.device))
-
-print(outputs)
-print(outputs.logits.shape)
-
-index.reset()

gpt4all/models/pythia_retro/__init__.py

@@ -1,2 +0,0 @@
-from .configuration_pythia_retro import PythiaRetroConfig
-from .modeling_pythia_retro import PythiaRetroForCausalLM

gpt4all/models/pythiaseek/__init__.py

@@ -1,2 +0,0 @@
-from .configuration_pythiaseek import PythiaSeekConfig
-from .modeling_pythiaseek import PythiaSeekForCausalLM

gpt4all/models/pythiaseek/configuration_pythiaseek.py

@@ -1,142 +0,0 @@
-# coding=utf-8
-# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPTNeoX model configuration"""
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "EleutherAI/gpt-neox-20b": "https://huggingface.co/EleutherAI/gpt-neox-20b/resolve/main/config.json",
-    # See all GPTNeoX models at https://huggingface.co/models?filter=gpt_neox
-}
-
-
-class PythiaSeekConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`GPTNeoXModel`]. It is used to instantiate an
-    GPTNeoX model according to the specified arguments, defining the model architecture. Instantiating a configuration
-    with the defaults will yield a similar configuration to that of the GPTNeoX
-    [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 50432):
-            Vocabulary size of the GPTNeoX model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`GPTNeoXModel`].
-        hidden_size (`int`, *optional*, defaults to 6144):
-            Dimension of the encoder layers and the pooler layer.
-        num_hidden_layers (`int`, *optional*, defaults to 44):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 64):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        intermediate_size (`int`, *optional*, defaults to 24576):
-            Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
-        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
-            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
-            `"relu"`, `"selu"` and `"gelu_new"` are supported.
-        rotary_pct (`float`, *optional*, defaults to 0.25):
-            percentage of hidden dimensions to allocate to rotary embeddings
-        rotary_emb_base (`int`, *optional*, defaults to 10000)
-            base for computing rotary embeddings frequency
-        classifier_dropout (`float`, *optional*, defaults to 0.1):
-            Argument used when doing token classification, used in the model [`GPTNeoXForTokenClassification`].
-
-            The dropout ratio for the hidden layer.
-        max_position_embeddings (`int`, *optional*, defaults to 2048):
-            The maximum sequence length that this model might ever be used with. Typically set this to something large
-            just in case (e.g., 512 or 1024 or 2048).
-        initializer_range (`float`, *optional*, defaults to 1e-5):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
-            The epsilon used by the layer normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models). Only
-            relevant if `config.is_decoder=True`.
-        use_parallel_residual (`bool`, *optional*, defaults to `True`):
-            Whether to use a "parallel" formulation in each Transformer layer, which can provide a slight training
-            speedup at large scales (e.g. 20B).
-        Example:
-
-    ```python
-    >>> from transformers import GPTNeoXConfig, GPTNeoXModel
-
-    >>> # Initializing a GPTNeoX gpt-neox-20b style configuration
-    >>> configuration = GPTNeoXConfig()
-
-    >>> # Initializing a model (with random weights) from the gpt-neox-20b style configuration
-    >>> model = GPTNeoXModel(configuration)  # doctest: +SKIP
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config  # doctest: +SKIP
-    ```"""
-    model_type = "gpt_neox"
-
-    def __init__(
-        self,
-        vocab_size=50432,
-        hidden_size=6144,
-        num_hidden_layers=44,
-        num_attention_heads=64,
-        intermediate_size=24576,
-        hidden_act="gelu",
-        rotary_pct=0.25,
-        rotary_emb_base=10000,
-        classifier_dropout=0.1,
-        max_position_embeddings=2048,
-        initializer_range=0.02,
-        layer_norm_eps=1e-5,
-        use_cache=True,
-        bos_token_id=0,
-        eos_token_id=2,
-        tie_word_embeddings=False,
-        use_parallel_residual=True,
-        encoder_dim=4096,
-        total_alpha_steps=0,
-        initial_alpha=1,
-        final_alpha=.5,
-        cross_attn_layer=9,
-        learnable_alpha=False,
-        **kwargs,
-    ):
-        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.intermediate_size = intermediate_size
-        self.hidden_act = hidden_act
-        self.rotary_pct = rotary_pct
-        self.rotary_emb_base = rotary_emb_base
-        self.classifier_dropout = classifier_dropout
-        self.initializer_range = initializer_range
-        self.layer_norm_eps = layer_norm_eps
-        self.use_cache = use_cache
-        self.tie_word_embeddings = tie_word_embeddings
-        self.use_parallel_residual = use_parallel_residual
-
-        self.encoder_dim = encoder_dim
-        self.total_alpha_steps = total_alpha_steps
-        self.initial_alpha = initial_alpha
-        self.final_alpha = final_alpha
-        # index of cross attention layer to add
-        self.cross_attn_layer = cross_attn_layer
-        self.learnable_alpha = learnable_alpha

gpt4all/models/pythiaseek/modeling_pythiaseek.py

@@ -1,899 +0,0 @@
-# coding=utf-8
-# Copyright 2022 EleutherAI The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch PythiaSeek model."""
-
-from typing import Optional, Tuple, Union
-
-import math
-import torch.nn.functional as F
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from transformers import AutoModel
-from transformers.activations import ACT2FN
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import logging
-from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
-from gpt4all.models.pythiaseek import PythiaSeekConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST = [
-    "EleutherAI/gpt-neox-20b",
-    # See all PythiaSeek models at https://huggingface.co/models?filter=gpt_neox
-]
-
-
-class PythiaSeekPreTrainedModel(PreTrainedModel):
-    """
-    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
-    models.
-    """
-
-    config_class = PythiaSeekConfig
-    base_model_prefix = "gpt_neox"
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["PythiaSeekLayer"]
-
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, PythiaSeekModel):
-            module.gradient_checkpointing = value
-
-
-class PythiaSeekAttention(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.num_attention_heads = config.num_attention_heads
-        self.hidden_size = config.hidden_size
-        self.head_size = self.hidden_size // self.num_attention_heads
-        self.rotary_ndims = int(self.head_size * config.rotary_pct)
-        max_positions = config.max_position_embeddings
-        self.register_buffer(
-            "bias",
-            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
-                1, 1, max_positions, max_positions
-            ),
-        )
-        self.register_buffer("masked_bias", torch.tensor(-1e9))
-        self.rotary_emb = RotaryEmbedding(
-            self.rotary_ndims, config.max_position_embeddings, base=config.rotary_emb_base
-        )
-        self.register_buffer(
-            "norm_factor",
-            torch.sqrt(torch.tensor(self.head_size, dtype=torch.float32)).to(torch.get_default_dtype()),
-            persistent=False,
-        )
-        self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size)
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-
-    def forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        attention_mask: torch.FloatTensor,
-        position_ids: torch.LongTensor,
-        head_mask: Optional[torch.FloatTensor] = None,
-        layer_past: Optional[Tuple[torch.Tensor]] = None,
-        use_cache: Optional[bool] = False,
-        output_attentions: Optional[bool] = False,
-    ):
-        has_layer_past = layer_past is not None
-
-        # Compute QKV
-        # Attention heads [batch, seq_len, hidden_size]
-        #   --> [batch, seq_len, (np * 3 * head_size)]
-        qkv = self.query_key_value(hidden_states)
-
-        # [batch, seq_len, (num_heads * 3 * head_size)]
-        #   --> [batch, seq_len, num_heads, 3 * head_size]
-        new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
-        qkv = qkv.view(*new_qkv_shape)
-
-        # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
-        query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
-        key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
-        value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)
-
-        # Compute rotary embeddings on rotary_ndims
-        query_rot = query[..., : self.rotary_ndims]
-        query_pass = query[..., self.rotary_ndims :]
-        key_rot = key[..., : self.rotary_ndims]
-        key_pass = key[..., self.rotary_ndims :]
-
-        # Compute token offset for rotary embeddings (when decoding)
-        seq_len = key.shape[-2]
-        if has_layer_past:
-            seq_len += layer_past[0].shape[-2]
-        cos, sin = self.rotary_emb(value, seq_len=seq_len)
-        query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
-        query = torch.cat((query, query_pass), dim=-1)
-        key = torch.cat((key, key_pass), dim=-1)
-
-        # Cache QKV values
-        if has_layer_past:
-            past_key = layer_past[0]
-            past_value = layer_past[1]
-            key = torch.cat((past_key, key), dim=-2)
-            value = torch.cat((past_value, value), dim=-2)
-        present = (key, value) if use_cache else None
-
-        # Compute attention
-        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
-        # Reshape outputs
-        attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size)
-        attn_output = self.dense(attn_output)
-
-        outputs = (attn_output, present)
-        if output_attentions:
-            outputs += (attn_weights,)
-
-        return outputs
-
-    @classmethod
-    def _split_heads(cls, tensor, num_attention_heads, attn_head_size):
-        """
-        Splits hidden dim into attn_head_size and num_attention_heads
-        """
-        # tensor: [bs, seq_len, hidden_size]
-        new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
-        # -> [bs, seq_len, num_attention_heads, attn_head_size]
-        tensor = tensor.view(new_shape)
-        # -> [bs, num_attention_heads, seq_len, attn_head_size]
-        tensor = tensor.permute(0, 2, 1, 3)
-        return tensor
-
-    @classmethod
-    def _merge_heads(cls, tensor, num_attention_heads, attn_head_size):
-        """
-        Merges attn_head_size dim and num_attn_heads dim into hidden dim
-        """
-        # tensor [bs, num_attention_heads, seq_len, attn_head_size]
-        tensor = tensor.permute(0, 2, 1, 3).contiguous()
-        # -> [bs, seq_len, num_attention_heads, attn_head_size]
-        tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size)
-        # -> [bs, seq_len, hidden_size]
-        return tensor
-
-    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
-        # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
-        # compute causal mask from causal mask buffer
-        batch_size, num_attention_heads, query_length, attn_head_size = query.size()
-        key_length = key.size(-2)
-
-        causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
-
-        query = query.view(batch_size * num_attention_heads, query_length, attn_head_size)
-        key = key.view(batch_size * num_attention_heads, key_length, attn_head_size)
-        attn_scores = torch.zeros(
-            batch_size * num_attention_heads,
-            query_length,
-            key_length,
-            dtype=query.dtype,
-            device=key.device,
-        )
-        attn_scores = torch.baddbmm(
-            attn_scores,
-            query,
-            key.transpose(1, 2),
-            beta=1.0,
-            alpha=(torch.tensor(1.0, dtype=self.norm_factor.dtype, device=self.norm_factor.device) / self.norm_factor),
-        )
-        attn_scores = attn_scores.view(batch_size, num_attention_heads, query_length, key_length)
-
-        mask_value = torch.finfo(attn_scores.dtype).min
-        # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
-        # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
-        mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(attn_scores.device)
-        attn_scores = torch.where(causal_mask, attn_scores, mask_value)
-
-        if attention_mask is not None:
-            # Apply the attention mask
-            attn_scores = attn_scores + attention_mask
-
-        attn_weights = nn.functional.softmax(attn_scores, dim=-1)
-        attn_weights = attn_weights.to(value.dtype)
-
-        # Mask heads if we want to
-        if head_mask is not None:
-            attn_weights = attn_weights * head_mask
-
-        attn_output = torch.matmul(attn_weights, value)
-        return attn_output, attn_weights
-
-class PythiaSeekCrossAttention(PythiaSeekAttention):
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.embed_dim = config.hidden_size
-        self.num_attention_heads = config.num_attention_heads
-        self.head_dim = self.embed_dim // self.num_attention_heads
-        if self.head_dim * self.num_attention_heads != self.embed_dim:
-            raise ValueError(
-                f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
-                f" `num_attention_heads`: {self.num_attention_heads})."
-            )
-        self.norm_factor = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype())
-
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-
-        
-    def _split_knn_attn_heads(self, tensor, num_attention_heads, attn_head_size):
-        new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
-        tensor = tensor.view(new_shape)
-
-        return tensor.permute(0, 1, 3, 2)
-
-        
-    def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
-        """
-        Merges attn_head_size dim and num_attn_heads dim into hidden dim
-        """
-        # tensor -> (bs, seq_len, num_attention_heads, head_dim)
-        tensor = tensor.permute(0, 2, 1, 3).contiguous()
-        new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
-        return tensor.view(new_shape)
-
-    def _attn(
-        self,
-        query,
-        key,
-        value,
-        attention_mask=None,
-        head_mask=None,
-    ):
-
-        # query -> (bs, num_attention_heads, seq_len, head_dim)
-        # key -> (bs, num_attention_heads, head_dim, neighbors)
-        # attn_weights -> (bs, num_attention_heads, seq_len, neighbors)
-        attn_weights = torch.matmul(query, key) / self.norm_factor
-
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-        attn_weights = attn_weights.to(value.dtype)
-
-        # Mask heads if we want to
-        if head_mask is not None:
-            attn_weights = attn_weights * head_mask
-
-        # value -> (bs, num_attention_heads, seq_len, head_dim)
-        # attn_weights -> (bs, num_attention_heads, seq_len, neighbors)
-        # attn_output -> (bs, num_attention_heads, seq_len, head_dim)
-        attn_output = torch.matmul(attn_weights, value)
-
-        return attn_output, attn_weights
-
-            
-    def forward(
-        self,
-        hidden_states: Optional[torch.FloatTensor],
-        encoder_hidden_states: Optional[torch.FloatTensor],
-        attention_mask: Optional[torch.FloatTensor] = None,
-        layer_past: Optional[Tuple[torch.Tensor]] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = False,
-        output_attentions: Optional[bool] = False,
-    ) -> Union[
-        Tuple[torch.Tensor, Tuple[torch.Tensor]],
-        Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
-    ]:
-
-        query = self.q_proj(hidden_states)
-        # if we are doing cross attention
-        key = self.k_proj(encoder_hidden_states)
-        value = self.v_proj(encoder_hidden_states)
-        # (bs, seq_len, dim) -> (bs, num_attention_heads, seq_len, head_dim)
-        query = self._split_heads(query, self.num_attention_heads, self.head_dim)
-        # (bs, dim) -> (bs, num_attention_heads,  head_dim)
-        key = self._split_knn_attn_heads(key, self.num_attention_heads, self.head_dim)
-        value = self._split_knn_attn_heads(value, self.num_attention_heads, self.head_dim)
-
-
-        value = value.permute(0, 3, 1, 2)
-        key = key.permute(0, 3, 2, 1)
-
-        if layer_past is not None:
-            past_key = layer_past[0]
-            past_value = layer_past[1]
-            key = torch.cat((past_key, key), dim=-2)
-            value = torch.cat((past_value, value), dim=-2)
-
-        if use_cache is True:
-            present = (key, value)
-        else:
-            present = None
-
-        # compute self-attention: V x Softmax(QK^T)
-        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
-        attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
-        attn_output = self.out_proj(attn_output)
-
-        outputs = (attn_output, present)
-        if output_attentions:
-            outputs += (attn_weights,)
-
-        return outputs  # a, present, (attentions)
-
-
-def attention_mask_func(attention_scores, ltor_mask):
-    attention_scores.masked_fill_(~ltor_mask, torch.finfo(attention_scores.dtype).min)
-    return attention_scores
-
-
-class RotaryEmbedding(torch.nn.Module):
-    def __init__(self, dim, max_position_embeddings, base=10000, device=None):
-        super().__init__()
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
-        self.register_buffer("inv_freq", inv_freq)
-
-        # Build here to make `torch.jit.trace` work.
-        self.max_seq_len_cached = max_position_embeddings
-        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
-        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
-        # Different from paper, but it uses a different permutation in order to obtain the same calculation
-        emb = torch.cat((freqs, freqs), dim=-1)
-        self.cos_cached = emb.cos()[None, None, :, :]
-        self.sin_cached = emb.sin()[None, None, :, :]
-
-    def forward(self, x, seq_len=None):
-        # x: [bs, num_attention_heads, seq_len, head_size]
-        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
-        if seq_len > self.max_seq_len_cached:
-            self.max_seq_len_cached = seq_len
-            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
-            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
-            # Different from paper, but it uses a different permutation in order to obtain the same calculation
-            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
-            self.cos_cached = emb.cos()[None, None, :, :]
-            self.sin_cached = emb.sin()[None, None, :, :]
-        return self.cos_cached[:seq_len, ...].to(x.device), self.sin_cached[:seq_len, ...].to(x.device)
-
-
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
-    gather_indices = position_ids[:, None, :, None]  # [bs, 1, seq_len, 1]
-    gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3])
-    cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
-    sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-class PythiaSeekMLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size)
-        self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.act = ACT2FN[config.hidden_act]
-
-    def forward(self, hidden_states):
-        hidden_states = self.dense_h_to_4h(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.dense_4h_to_h(hidden_states)
-        return hidden_states
-
-
-class PythiaSeekLayer(nn.Module):
-    def __init__(self, config, add_cross_attention=False):
-        super().__init__()
-        self.use_parallel_residual = config.use_parallel_residual
-        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.attention = PythiaSeekAttention(config)
-        self.mlp = PythiaSeekMLP(config)
-
-        self.add_cross_attention = add_cross_attention
-
-        if add_cross_attention:
-            self.cross_attn_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-            self.cross_attn = PythiaSeekCrossAttention(config)
-            self.cross_attn_mlp = PythiaSeekMLP(config)
-            
-            self.total_alpha_steps = config.total_alpha_steps
-            self.initial_alpha = config.initial_alpha
-            self.final_alpha = config.final_alpha
-            self.learnable_alpha = config.learnable_alpha
-            if self.learnable_alpha:
-                self.alpha = nn.Parameter(torch.zeros(1), requires_grad=True)
-
-    def forward(
-        self,
-        hidden_states: Optional[torch.FloatTensor],
-        encoder_hidden_states: Optional[torch.FloatTensor],
-        attention_mask: Optional[torch.FloatTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = False,
-        layer_past: Optional[Tuple[torch.Tensor]] = None,
-        output_attentions: Optional[bool] = False,
-        step: Optional[int] = None,
-    ):
-        ln_hidden_states = self.input_layernorm(hidden_states)
-        attention_layer_outputs = self.attention(
-            ln_hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            layer_past=layer_past,
-            head_mask=head_mask,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-        )
-        attn_output = attention_layer_outputs[0]  # output_attn: attn_output, present, (attn_weights)
-        outputs = attention_layer_outputs[1:]
-
-        if self.use_parallel_residual:
-            # pseudocode:
-            # x = x + attn(ln1(x)) + mlp(ln2(x))
-            mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))
-            self_attention_residual = mlp_output + attn_output + hidden_states
-        else:
-            # pseudocode:
-            # x = x + attn(ln1(x))
-            # x = x + mlp(ln2(x))
-            attn_output = attn_output + hidden_states
-            mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
-            self_attention_residual = mlp_output + attn_output
-
-        # encoder_hidden_states -> (bs, knn, encoder_dim)
-        if encoder_hidden_states.dtype != ln_hidden_states.dtype:
-            encoder_hidden_states = encoder_hidden_states.to(ln_hidden_states.dtype)
-
-        if self.add_cross_attention:
-            encoder_normed = self.cross_attn_ln(encoder_hidden_states)
-
-            # cross_attn_outputs -> (bs, seq_len, dim)
-            cross_attn_output = self.cross_attn(
-                ln_hidden_states,
-                encoder_hidden_states=encoder_normed,
-                attention_mask=attention_mask,
-                head_mask=head_mask,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-            )
-
-            cross_attn_ff = self.cross_attn_mlp(
-                cross_attn_output[0] 
-            )
-
-            if self.learnable_alpha:
-                alpha = F.sigmoid(self.alpha)
-            elif step is not None:
-                alpha = self._update_alpha(step)
-            else:
-                alpha = 0.5 
-
-            hidden_states = (1 - alpha) * cross_attn_ff +  alpha * self_attention_residual  
-        else:
-            hidden_states = self_attention_residual
-
-        if use_cache:
-            outputs = (hidden_states,) + outputs
-        else:
-            outputs = (hidden_states,) + outputs[1:]
-
-        return outputs  # hidden_states, present, (attentions)
-        
-    def _update_alpha(self, current_step):
-        """
-        Computes the learning rate for the current step using a cosine decay schedule.
-
-        Args:
-            initial_lr (float): The initial learning rate.
-            final_lr (float): The final learning rate.
-            total_steps (int): The total number of steps in the schedule.
-            current_step (int): The current step.
-
-        Returns:
-            float: The learning rate for the current step.
-        """
-        initial_alpha = 1
-        final_alpha = .5
-        if current_step >= self.total_alpha_steps:
-            return final_alpha
-
-        # Compute the cosine decay factor
-        cosine_decay = 0.5 * (1 + math.cos(math.pi * current_step / self.total_alpha_steps))
-
-        # Compute the current learning rate
-        alpha = final_alpha + (initial_alpha - final_alpha) * cosine_decay
-
-        return alpha
-
-
-
-class PythiaSeekModel(PythiaSeekPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.config = config
-
-        self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
-
-        if isinstance(config.cross_attn_layer, int):
-            cross_attn_layers = [config.cross_attn_layer]
-        # if we add cross attention to all layers
-        elif config.cross_attn_layer is None:
-            cross_attn_layers = list(range(config.num_hidden_layers))
-        else:
-            cross_attn_layers = config.cross_attn_layer
-
-        self.layers = nn.ModuleList([PythiaSeekLayer(config, add_cross_attention=i in cross_attn_layers) for i in range(config.num_hidden_layers)])
-        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
-        self.gradient_checkpointing = False
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_in
-
-    def set_input_embeddings(self, value):
-        self.embed_in = value
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        encoder_hidden_states: Optional[torch.FloatTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        step: Optional[int] = None
-    ) -> Union[Tuple, BaseModelOutputWithPast]:
-        r"""
-        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
-            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
-            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
-            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
-        use_cache (`bool`, *optional*):
-            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
-            `past_key_values`).
-        """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
-        elif input_ids is not None:
-            input_shape = input_ids.size()
-        elif inputs_embeds is not None:
-            input_shape = inputs_embeds.size()[:-1]
-        else:
-            raise ValueError("You have to specify either input_ids or inputs_embeds")
-
-        batch_size, seq_length = input_shape
-
-        if past_key_values is None:
-            past_length = 0
-            past_key_values = tuple([None] * self.config.num_hidden_layers)
-        else:
-            past_length = past_key_values[0][0].size(-2)
-
-        if position_ids is None:
-            device = input_ids.device if input_ids is not None else inputs_embeds.device
-            position_ids = torch.arange(past_length, seq_length + past_length, dtype=torch.long, device=device)
-            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
-        else:
-            position_ids = position_ids.view(-1, seq_length).long()
-
-        # Attention mask.
-        if attention_mask is not None:
-            assert batch_size > 0, "batch_size has to be defined and > 0"
-            attention_mask = attention_mask.view(batch_size, -1)
-            # We create a 3D attention mask from a 2D tensor mask.
-            # Sizes are [batch_size, 1, 1, to_seq_length]
-            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
-            # this attention mask is more simple than the triangular masking of causal attention
-            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
-            attention_mask = attention_mask[:, None, None, :]
-
-            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
-            # masked positions, this operation will create a tensor which is 0.0 for
-            # positions we want to attend and the dtype's smallest value for masked positions.
-            # Since we are adding it to the raw scores before the softmax, this is
-            # effectively the same as removing these entirely.
-            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
-            attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
-        # Prepare head mask if needed
-        # 1.0 in head_mask indicate we keep the head
-        # attention_probs has shape bsz x n_heads x N x N
-        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
-        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
-        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_in(input_ids)
-
-        hidden_states = inputs_embeds
-
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
-        presents = () if use_cache else None
-        all_attentions = () if output_attentions else None
-        all_hidden_states = () if output_hidden_states else None
-        for i, (layer, layer_past) in enumerate(zip(self.layers, past_key_values)):
-            if output_hidden_states:
-                all_hidden_states = all_hidden_states + (hidden_states,)
-
-            if self.gradient_checkpointing and self.training:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for layer_past
-                        return module(*inputs, use_cache, None, output_attentions, step)
-
-                    return custom_forward
-
-                outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(layer),
-                    hidden_states,
-                    encoder_hidden_states,
-                    attention_mask,
-                    position_ids,
-                    head_mask[i],
-                )
-            else:
-                outputs = layer(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    head_mask=head_mask[i],
-                    layer_past=layer_past,
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                    step=step,
-                )
-            hidden_states = outputs[0]
-            if use_cache is True:
-                presents = presents + (outputs[1],)
-            if output_attentions:
-                all_attentions = all_attentions + (outputs[2 if use_cache else 1],)
-
-        hidden_states = self.final_layer_norm(hidden_states)
-        # Add last hidden state
-        if output_hidden_states:
-            all_hidden_states = all_hidden_states + (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)
-
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=presents,
-            hidden_states=all_hidden_states,
-            attentions=all_attentions,
-        )
-
-
-class PythiaSeekForCausalLM(PythiaSeekPreTrainedModel):
-    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.gpt_neox = PythiaSeekModel(config)
-        self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-        self.hidden_size = config.hidden_size
-        self.encoder_dim = config.encoder_dim
-
-        if self.hidden_size != self.encoder_dim:
-            self.enc_dec_proj = nn.Sequential(nn.Linear(config.encoder_dim, config.hidden_size * 4),
-                                              nn.ReLU(),
-                                              nn.Linear(config.hidden_size * 4, config.hidden_size),
-                                              )
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_output_embeddings(self):
-        return self.embed_out
-
-    def set_output_embeddings(self, new_embeddings):
-        self.embed_out = new_embeddings
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        encoder_hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        step: Optional[int] = None
-    ) -> Union[Tuple, CausalLMOutputWithPast]:
-        r"""
-        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
-            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
-            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
-            `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional tensors are
-            only required when the model is used as a decoder in a Sequence to Sequence model.
-
-            Contains pre-computed hidden-states (key and values in the self-attention blocks that can be used (see
-            `past_key_values` input) to speed up sequential decoding.
-
-            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
-            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
-            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
-            `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
-            ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
-        use_cache (`bool`, *optional*):
-            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
-            `past_key_values`).
-
-        Returns:
-
-        Example:
-
-        ```python
-        >>> from transformers import AutoTokenizer, PythiaSeekForCausalLM, PythiaSeekConfig
-        >>> import torch
-
-        >>> tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
-        >>> config = PythiaSeekConfig.from_pretrained("EleutherAI/gpt-neox-20b")
-        >>> config.is_decoder = True
-        >>> model = PythiaSeekForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b", config=config)
-
-        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
-        >>> outputs = model(**inputs)
-
-        >>> prediction_logits = outputs.logits
-        ```"""
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if self.hidden_size != self.encoder_dim:
-            encoder_hidden_states = encoder_hidden_states.to(self.enc_dec_proj[0].weight.dtype)
-            encoder_hidden_states = self.enc_dec_proj(encoder_hidden_states)
-
-        outputs = self.gpt_neox(
-            input_ids,
-            encoder_hidden_states=encoder_hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            past_key_values=past_key_values,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            step=step,
-        )
-
-        hidden_states = outputs[0]
-        lm_logits = self.embed_out(hidden_states)
-
-        lm_loss = None
-        if labels is not None:
-            # move labels to correct device to enable model parallelism
-            labels = labels.to(lm_logits.device)
-            # we are doing next-token prediction; shift prediction scores and input ids by one
-            shift_logits = lm_logits[:, :-1, :].contiguous()
-            labels = labels[:, 1:].contiguous()
-            loss_fct = CrossEntropyLoss()
-            lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
-
-        if not return_dict:
-            output = (lm_logits,) + outputs[1:]
-            return ((lm_loss,) + output) if lm_loss is not None else output
-
-        return CausalLMOutputWithPast(
-            loss=lm_loss,
-            logits=lm_logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
-    ):
-        input_shape = input_ids.shape
-
-        # cut decoder_input_ids if past is used
-        if past_key_values and past_key_values[0] is not None:
-            input_ids = input_ids[:, -1:]
-
-        position_ids = kwargs.get("position_ids", None)
-        if attention_mask is not None and position_ids is None:
-            # create position_ids on the fly for batch generation
-            position_ids = attention_mask.long().cumsum(-1) - 1
-            position_ids.masked_fill_(attention_mask == 0, 1)
-            if past_key_values:
-                position_ids = position_ids[:, -1].unsqueeze(-1)
-
-        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
-        if attention_mask is None:
-            attention_mask = input_ids.new_ones(input_shape)
-
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and past_key_values is None:
-            model_inputs = {"inputs_embeds": inputs_embeds}
-        else:
-            model_inputs = {"input_ids": input_ids}
-
-        model_inputs.update(
-            {
-                "attention_mask": attention_mask,
-                "past_key_values": past_key_values,
-                "position_ids": position_ids,
-            }
-        )
-
-        return model_inputs
-
-    def _reorder_cache(self, past_key_values, beam_idx):
-        reordered_past = ()
-        for layer_past in past_key_values:
-            reordered_past += (
-                tuple(past_state.index_select(0, beam_idx) for past_state in layer_past[:2]) + layer_past[2:],
-            )
-        return reordered_past

gpt4all/models/pythiaseek/test_pythiaseek.py

@@ -1,55 +0,0 @@
-import torch
-from gpt4all.models import PythiaSeekConfig, PythiaSeekForCausalLM
-from transformers import AutoTokenizer, AutoModel
-
-# seed torch 
-
-torch.manual_seed(0)
-
-config = PythiaSeekConfig(encoder_dim=384, 
-                          hidden_size=1024, 
-                          intermediate_size=4096,
-                          num_hidden_layers=4,
-                          num_attention_heads=4)
-print("loaded config")
-
-print("loading model")
-model = PythiaSeekForCausalLM(config)
-print("loaded model")
-
-
-tokenizer =  AutoTokenizer.from_pretrained("EleutherAI/pythia-1b")
-tokenizer.pad_token = tokenizer.eos_token
-tokenizer.model_max_length = 1024
-
-encoder_tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-encoder = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-
-
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-
-text = "The quick brown fox jumps over the lazy dog."
-print("Encoded knn")
-tokenized = encoder_tokenizer(text, return_tensors="pt")
-
-# bs, seq_len, dim
-encodings = mean_pooling(encoder(**tokenized), tokenized["attention_mask"])
-
-# make 2 neighbors
-# (bs, knn, encoding_dim)
-encoder_hidden_states = torch.stack([encodings, encodings]).squeeze().unsqueeze(0)
-
-inputs = "What did the fox do?"
-
-print("Encoded inputs")
-tokenized_input = tokenizer([inputs], padding="max_length", truncation=True, return_tensors="pt")
-
-print("Running model")
-outputs = model(**tokenized_input, encoder_hidden_states=encoder_hidden_states)
-
-print(outputs)
-print(outputs[0].shape)
-

gpt4all/train/metrics.py

@@ -1,50 +0,0 @@
-from collections import Counter
-import string 
-import re
-
-
-# adapted from huggingface
-def normalize_answer(s):
-    """Lower text and remove punctuation, articles and extra whitespace."""
-
-    def remove_articles(text):
-        return re.sub(r"\b(a|an|the)\b", " ", text)
-
-    def white_space_fix(text):
-        return " ".join(text.split())
-
-    def remove_punc(text):
-        exclude = set(string.punctuation)
-        return "".join(ch for ch in text if ch not in exclude)
-
-    def lower(text):
-        return text.lower()
-
-    return white_space_fix(remove_articles(remove_punc(lower(s))))
-
-
-def f1_score(predictions, ground_truths):
-    total_f1 = []
-    for prediction, ground_truth in zip(predictions, ground_truths):
-        prediction_tokens = normalize_answer(prediction).split()
-        ground_truth_tokens = normalize_answer(ground_truth).split()
-        common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
-        num_same = sum(common.values())
-        if num_same == 0:
-            total_f1.append(0)
-            continue
-
-        precision = 1.0 * num_same / len(prediction_tokens)
-        recall = 1.0 * num_same / len(ground_truth_tokens)
-        f1 = (2 * precision * recall) / (precision + recall)
-        total_f1.append(f1)
-
-    return total_f1
-
-
-def exact_match_score(predictions, ground_truths):
-    exact_scores = []
-    for prediction, ground_truth in zip(predictions, ground_truths):
-        exact_scores.append(normalize_answer(prediction) == normalize_answer(ground_truth))
-
-    return exact_scores

gpt4all/train/train_mem_retrieval.py

@@ -1,328 +0,0 @@
-import os
-import torch.nn.functional as F
-from transformers import AutoTokenizer, get_scheduler, AutoConfig
-import torch
-from torch.optim import AdamW
-from argparse import ArgumentParser
-from gpt4all.utils.read import read_config
-from accelerate import Accelerator
-from accelerate.utils import DummyScheduler, DummyOptim, set_seed
-from gpt4all.data.retrieval_dataloader import load_memory_augmented_data
-from torchmetrics import MeanMetric
-from tqdm import tqdm
-from gpt4all.models import LetheForCausalLM
-from gpt4all.models.lethe.modeling_lethe import MemoryIndex
-import wandb
-import pyarrow as pa
-from pyarrow import feather
-
-torch.backends.cuda.matmul.allow_tf32 = True
-
-def format_metrics(metrics, split, prefix=""):
-    log = f"[{split}]" + prefix
-    log += " ".join([f"{key}: {value:.4f}" for key, value in metrics.items()])
-
-    return log
-
-
-def calculate_per_example_loss(logits, labels):
-    lm_logits = logits[:, :-1, :].contiguous()
-    lm_labels = labels[:, 1:].contiguous()
-    loss = F.cross_entropy(lm_logits.view(-1, lm_logits.size(-1)), lm_labels.view(-1), reduction="none")
-    loss = loss.reshape(labels.shape[0], -1).mean(dim=-1)
-
-    # return tensor of shape (B,) where B is the batch size
-    return loss.cpu().tolist()
-
-
-def evaluate(model, index, config, val_dataloader, main_process=False):
-    model.eval()
-    val_loss = MeanMetric(nan_strategy="error").to(model.device)
-
-    ids = []
-    losses = []
-    with torch.no_grad():
-        for batch in tqdm(val_dataloader, disable=not main_process):
-            batch["id"] = batch["id"].detach().cpu()
-            memories = batch["retrieved_context"]
-
-            memories = memories[:, :config["num_neighbors_to_store"], :]
-            memories = memories.reshape(-1, memories.shape[-1])
-
-            for chunk_start in range(0, memories.shape[0], config["mem_chunk_size"]):
-                chunk_end = min(memories.shape[0], chunk_start + config["mem_chunk_size"])
-                mem_chunk = memories[chunk_start:chunk_end]
-                model(input_ids=mem_chunk)
-
-            qa_inputs = batch["input_ids"]
-            qa_labels = batch["labels"]
-            outputs = model(input_ids=qa_inputs, 
-                            labels=qa_labels, 
-            )
-
-            del memories
-            torch.cuda.empty_cache()
-            loss_values = accelerator.gather_for_metrics({"loss": outputs["loss"].detach()})
-
-            index.reset()
-            val_loss.update(loss_values["loss"])
-
-            per_example_loss = calculate_per_example_loss(outputs["logits"], qa_labels)
-
-            losses.extend(per_example_loss)
-            ids.extend(batch["id"].tolist())
-
-    ids = pa.array(ids)
-    losses = pa.array(losses)
-    schema = pa.schema([("loss", pa.float64()), ("id", pa.int32())])
-    table = pa.Table.from_arrays([losses, ids], schema=schema)
-    return val_loss, table
-
-
-def train(accelerator, config):
-    set_seed(config['seed'])
-
-    accelerator.print(config)
-    accelerator.print(f"Using {accelerator.num_processes} GPUs")
-
-    tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'], model_max_length=config['max_length'])
-    # if no pad token, set it to eos
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-        
-    with accelerator.main_process_first():
-        train_dataloader, val_dataloader = load_memory_augmented_data(config, tokenizer) 
-
-
-    if accelerator.state.deepspeed_plugin is not None:
-        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
-            "gradient_accumulation_steps"
-        ]
-
-    accelerator.print(f"Len of train_dataloader: {len(train_dataloader)}")
-    total_num_steps = (len(train_dataloader) / gradient_accumulation_steps) * config["num_epochs"]
-    # instead of decaying to zero, decay to ratio of min_lr / lr
-    accelerator.print(f"Total training steps: {total_num_steps}")
-
-    checkpoint = config["gradient_checkpointing"]
-    
-    model_config = AutoConfig.from_pretrained(config["model_name"])
-
-    head_size = model_config.hidden_size // model_config.num_attention_heads
-    index = MemoryIndex(head_size,
-                        config["num_memories_per_index"],
-                        model_config.num_attention_heads
-    )
-    model = LetheForCausalLM.from_pretrained(config["model_name"], 
-                                             revision=config['version'] if 'version' in config else None,
-                                             use_cache=False if checkpoint else True,
-                                             memory_attn_layer=config["memory_attn_layer"],
-                                             num_neighbors_to_retrieve=config["num_neighbors_to_retrieve"],
-                                             index=index,
-                                             tracker=accelerator.get_tracker("wandb"),
-                                            )
-
-
-    accelerator.print(f"Training a {model.num_parameters():,} parameter model")
-    if checkpoint:
-        model.gradient_checkpointing_enable()
-
-    optimizer_cls = (
-        AdamW
-        if accelerator.state.deepspeed_plugin is None
-        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
-        else DummyOptim
-    )
-
-    # karpathy doesn't decay embeddding, maybe we should exclude
-    # https://github.com/karpathy/minGPT/commit/bbbdac74fa9b2e55574d70056163ffbae42310c1#diff-2075fa9c224b395be5bda85544dd36572b59c76c54562819eadadbf268602834R157s
-    optimizer = optimizer_cls(model.parameters(), lr=config["lr"], weight_decay=config["weight_decay"])
-
-
-    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
-    if config["scheduler"] or "scheduler" in accelerator.state.deepspeed_plugin.deepspeed_config:
-        if (
-            accelerator.state.deepspeed_plugin is None
-            or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
-        ):
-            scheduler = get_scheduler(
-                name="cosine",
-                optimizer=optimizer,
-                num_warmup_steps=config["warmup_steps"] * accelerator.num_processes,
-                num_training_steps=total_num_steps,
-            )
-        else:
-            scheduler = DummyScheduler(
-                optimizer, total_num_steps=total_num_steps, warmup_num_steps=config["warmup_steps"]
-            )
-        model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
-                model, optimizer, train_dataloader, val_dataloader, scheduler
-        )
-        use_scheduler = True
-    else:
-        model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(
-                model, optimizer, train_dataloader, val_dataloader
-        )
-        use_scheduler = False
-
-    #  setup for saving training states in case preemption
-    if use_scheduler:
-        accelerator.register_for_checkpointing(scheduler)
-
-    if config["checkpoint"]:
-        accelerator.load_state(config["checkpoint"])
-        accelerator.print(f"Resumed from checkpoint: {config['checkpoint']}")
-        path = os.path.basename(config["train_args"]["resume_from_checkpoint"])
-        training_difference = os.path.splitext(path)[0]
-        resume_step = int(training_difference.replace("step_", ""))
-        accelerator.skip_first_batches(train_dataloader, resume_step)
-        accelerator.print(f"Resuming from step {resume_step}")
-
-    # log gradients
-    if accelerator.is_main_process and config["wandb"]:
-        wandb.watch(model, log_freq=config["log_grads_every"], log="all")
-
-    main_process = accelerator.is_main_process
-
-    for epoch in range(config["num_epochs"]):
-        train_loss = MeanMetric(nan_strategy="error").to(model.device)
-        for step, batch in enumerate(tqdm(train_dataloader, disable=not main_process)):
-            curr_step = step + epoch * len(train_dataloader)
-            memories = batch["retrieved_context"]
-            memories = memories[:, :config["num_neighbors_to_store"], :]
-            memories = memories.reshape(-1, memories.shape[-1])
-
-            # need to set to eval so we don't do mem attn as it's slow
-            model.eval()
-            with torch.no_grad():
-                for chunk_start in range(0, memories.shape[0], config["mem_chunk_size"]):
-                    chunk_end = min(memories.shape[0], chunk_start + config["mem_chunk_size"])
-                    mem_chunk = memories[chunk_start:chunk_end]
-                    model(input_ids=mem_chunk)
-
-            del memories
-            torch.cuda.empty_cache()
-
-            model.train()
-            qa_inputs = batch["input_ids"]
-            qa_labels = batch["labels"]
-            outputs = model(input_ids=qa_inputs, 
-                            labels=qa_labels, 
-                            log_attn_scores=True,
-                            step=curr_step,
-                            save_kv=False,
-            )
-            loss = outputs.loss
-            if config["wandb"]:
-                accelerator.log({"loss": loss}, step=curr_step)
-
-
-            # gather loss before backprop in case of gradient accumulation
-            loss_values = accelerator.gather_for_metrics({"loss": loss.detach().float()})
-            train_loss.update(loss_values["loss"])
-
-            loss = loss / gradient_accumulation_steps
-            accelerator.backward(loss)
-            # !! don't reset index until after backwards pass
-            index.reset()
-            # get gradient norm of all params
-
-            # log LR in case something weird happens 
-            if config["wandb"]:
-                if step > 0 and step % (config["log_lr_every"] ) == 0:
-                    lr = optimizer.param_groups[0]["lr"]
-                    accelerator.log({"lr": lr}, step=curr_step)
-
-            if (step + 1) % gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
-                optimizer.step()
-                if use_scheduler:
-                    scheduler.step()
-                optimizer.zero_grad()
-
-            if step > 0 and config["save_every"] > 0 and step % config["save_every"] == 0:
-                # accelerator.save_state(f"{config['output_dir']}/step_{curr_step}")
-                unwrapped_model = accelerator.unwrap_model(model)
-                unwrapped_model.save_pretrained(
-                        f"{config['output_dir']}/step_{step}",
-                        is_main_process=accelerator.is_main_process,
-                        save_function=accelerator.save,
-                        state_dict=accelerator.get_state_dict(model),
-                )
-
-            if step > 0 and (step % config["eval_every"] == 0 or step == len(train_dataloader) - 1):
-                val_loss, loss_table = evaluate(model, index, config, val_dataloader, main_process=main_process)
-
-                local_rank = accelerator.process_index
-                feather.write_feather(loss_table, f"{config['output_dir']}/val_losses_step_{curr_step}_rank_{local_rank}.feather")
-
-                log_train = {
-                        "train_loss": train_loss.compute()
-                    }
-                log_val = {
-                    "val_loss": val_loss.compute(),
-                }
-
-                if config["wandb"]:
-                    curr_step = step + epoch * len(train_dataloader)
-                    accelerator.log({**log_train, **log_val}, step=curr_step)
-
-                accelerator.print(f"Current LR: {optimizer.param_groups[0]['lr']}")
-                accelerator.print(format_metrics(log_train, "train", f" step {step} "))
-                accelerator.print(format_metrics(log_val, "val", f" step {step} "))
-
-                train_loss.reset()
-
-        accelerator.print(f"Epoch {epoch} finished")
-        accelerator.wait_for_everyone()
-        unwrapped_model = accelerator.unwrap_model(model)
-        if config["push_to_hub"]:
-            accelerator.print(f"Pushing to HF hub")
-            try:
-                if accelerator.is_main_process:
-                    unwrapped_model.push_to_hub(config["save_name"] + f"-epoch_{epoch}", private=True)
-
-            except Exception as e:
-                accelerator.print(e)
-                accelerator.print(f"Failed to push to hub")
-
-        unwrapped_model.save_pretrained(
-            f"{config['output_dir']}/epoch_{epoch}",
-            is_main_process=accelerator.is_main_process,
-            save_function=accelerator.save,
-            state_dict=accelerator.get_state_dict(model),
-        )
-            
-    accelerator.wait_for_everyone()
-    unwrapped_model = accelerator.unwrap_model(model)
-    unwrapped_model.save_pretrained(
-        f"{config['output_dir']}/final",
-        is_main_process=accelerator.is_main_process,
-        save_function=accelerator.save,
-        state_dict=accelerator.get_state_dict(model),
-    )
-
-    accelerator.end_training()
-
-    
-
-if __name__ == "__main__":
-    # parse arguments by reading in a config
-    parser = ArgumentParser()
-    parser.add_argument("--config", type=str, default="config.yaml")
-
-    args = parser.parse_args()
-
-    config = read_config(args.config)
-
-    if config["wandb"]:
-        accelerator = Accelerator(log_with="wandb")
-        accelerator.init_trackers(
-            project_name=config["wandb_project_name"],
-            config=config,
-            init_kwargs={"wandb": {"entity": config["wandb_entity"]}},
-        )
-    else:
-        accelerator = Accelerator()
-
-    train(accelerator, config=config)

gpt4all/train/train_retrieval.py

@@ -1,307 +0,0 @@
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer, get_scheduler
-import torch
-from torch.optim import AdamW
-from argparse import ArgumentParser
-from gpt4all.utils.read import read_config
-from accelerate import Accelerator
-from accelerate.utils import DummyScheduler, DummyOptim, set_seed
-from peft import get_peft_model, LoraConfig, TaskType
-from gpt4all.data.retrieval_dataloader import load_retrieval_augmented_data
-from torchmetrics import MeanMetric
-from tqdm import tqdm
-from gpt4all.models import GPTJRForCausalLM, PythiaSeekForCausalLM
-import wandb
-
-torch.backends.cuda.matmul.allow_tf32 = True
-
-def format_metrics(metrics, split, prefix=""):
-    log = f"[{split}]" + prefix
-    log += " ".join([f"{key}: {value:.4f}" for key, value in metrics.items()])
-
-    return log
-
-
-def evaluate(model, val_dataloader, step, main_process=False):
-    model.eval()
-    val_loss = MeanMetric(nan_strategy="error").to(model.device)
-
-    with torch.no_grad():
-        for batch in tqdm(val_dataloader, disable=not main_process):
-            outputs = model(input_ids=batch["input_ids"], 
-                            labels=batch["labels"], 
-                            encoder_hidden_states=batch["encoder_hidden_states"],
-                            step=step)
-            loss_values = accelerator.gather_for_metrics({"loss": outputs["loss"].detach()})
-
-
-            val_loss.update(loss_values["loss"])
-
-    return val_loss
-
-
-def train(accelerator, config):
-    set_seed(config['seed'])
-
-    accelerator.print(config)
-    accelerator.print(f"Using {accelerator.num_processes} GPUs")
-
-    tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'], model_max_length=config['max_length'])
-    # if no pad token, set it to eos
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-        
-    with accelerator.main_process_first():
-        train_dataloader, val_dataloader = load_retrieval_augmented_data(config, tokenizer) 
-
-    if accelerator.state.deepspeed_plugin is not None:
-        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
-            "gradient_accumulation_steps"
-        ]
-
-    accelerator.print(f"Len of train_dataloader: {len(train_dataloader)}")
-    total_num_steps = (len(train_dataloader) / gradient_accumulation_steps) * config["num_epochs"]
-    # instead of decaying to zero, decay to ratio of min_lr / lr
-    accelerator.print(f"Total training steps: {total_num_steps}")
-
-    checkpoint = config["gradient_checkpointing"]
-    #ensures back compat with non retrieval models
-    if 'encoder_dim' in config:
-        if "gptj" in config["model_name"]:
-            model = GPTJRForCausalLM.from_pretrained(config["model_name"], 
-                                                    revision=config['version'] if 'version' in config else None,
-                                                    use_cache=False if checkpoint else True,
-                                                    encoder_dim=config["encoder_dim"],
-                                                    # hardcoded!! TODO: change to config
-                                                    total_alpha_steps=10000,
-                                                    ) 
-        elif "pythia" in config["model_name"]:
-            cross_attn_layer = config["cross_attn_layer"]
-            learnable_alpha = config["learnable_alpha"]
-            model, info = PythiaSeekForCausalLM.from_pretrained(config["model_name"], 
-                                                    revision=config['version'] if 'version' in config else None,
-                                                    use_cache=False if checkpoint else True,
-                                                    encoder_dim=config["encoder_dim"],
-                                                    # hardcoded!! TODO: change to config
-                                                    total_alpha_steps=2500,
-                                                    # mem transformer uses 9/12, pythia1-b is 16 layers
-                                                    cross_attn_layer=cross_attn_layer,
-                                                    learnable_alpha=learnable_alpha,
-                                                    output_loading_info=True
-                                                    )
-            # freeze all pretrained layers
-            if config["freeze_pretrained"]:
-                for name, param in model.named_parameters():
-                    if name not in info["missing_keys"]:
-                        param.requires_grad = False
-
-    else:
-        model = AutoModelForCausalLM.from_pretrained(config["model_name"], 
-                                                     use_cache=False if checkpoint else True,
-                                                     trust_remote_code=True) 
-
-    accelerator.print(f"Training a {model.num_parameters():,} parameter model")
-    if checkpoint:
-        model.gradient_checkpointing_enable()
-
-    if config["lora"]:
-        peft_config = LoraConfig(
-            # should R be configurable?
-            task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
-        )
-        model = get_peft_model(model, peft_config)
-        model.print_trainable_parameters()
-
-    optimizer_cls = (
-        AdamW
-        if accelerator.state.deepspeed_plugin is None
-        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
-        else DummyOptim
-    )
-
-    # karpathy doesn't decay embeddding, maybe we should exclude
-    # https://github.com/karpathy/minGPT/commit/bbbdac74fa9b2e55574d70056163ffbae42310c1#diff-2075fa9c224b395be5bda85544dd36572b59c76c54562819eadadbf268602834R157s
-    optimizer = optimizer_cls(model.parameters(), lr=config["lr"], weight_decay=config["weight_decay"])
-
-
-    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
-    if config["scheduler"] or "scheduler" in accelerator.state.deepspeed_plugin.deepspeed_config:
-        if (
-            accelerator.state.deepspeed_plugin is None
-            or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
-        ):
-            scheduler = get_scheduler(
-                name="cosine",
-                optimizer=optimizer,
-                num_warmup_steps=config["warmup_steps"] * accelerator.num_processes,
-                num_training_steps=total_num_steps,
-            )
-        else:
-            scheduler = DummyScheduler(
-                optimizer, total_num_steps=total_num_steps, warmup_num_steps=config["warmup_steps"]
-            )
-        model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
-                model, optimizer, train_dataloader, val_dataloader, scheduler
-        )
-        scheduler = True
-    else:
-        model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(
-                model, optimizer, train_dataloader, val_dataloader
-        )
-        scheduler = False
-
-
-    #  setup for saving training states in case preemption
-    if scheduler:
-        accelerator.register_for_checkpointing(scheduler)
-
-    if config["checkpoint"]:
-        accelerator.load_state(config["checkpoint"])
-        accelerator.print(f"Resumed from checkpoint: {config['checkpoint']}")
-        path = os.path.basename(config["train_args"]["resume_from_checkpoint"])
-        training_difference = os.path.splitext(path)[0]
-        resume_step = int(training_difference.replace("step_", ""))
-        accelerator.skip_first_batches(train_dataloader, resume_step)
-        accelerator.print(f"Resuming from step {resume_step}")
-
-
-    # log gradients
-    if accelerator.is_main_process and config["wandb"]:
-        wandb.watch(model, log_freq=config["log_grads_every"], log="all")
-
-    main_process = accelerator.is_main_process
-
-    if learnable_alpha is False and isinstance(cross_attn_layer, int):
-        accelerator.log({"alpha": model.gpt_neox.layers[cross_attn_layer]._update_alpha(0)}, step=0)
-
-    for epoch in range(config["num_epochs"]):
-        train_loss = MeanMetric(nan_strategy="error").to(model.device)
-        for step, batch in enumerate(tqdm(train_dataloader, disable=not main_process)):
-            model.train()
-            curr_step = step + len(train_dataloader) * epoch
-            outputs = model(input_ids=batch["input_ids"], 
-                            labels=batch["labels"], 
-                            encoder_hidden_states=batch["encoder_hidden_states"],
-                            step=curr_step)
-            loss = outputs.loss
-
-            if config["debug"]:
-                logits = outputs.logits
-                pred_tokens = torch.argmax(logits, dim=-1)
-                labels = batch["labels"].clone()
-
-                for row in range(labels.shape[0]):
-                    curr_label = labels[row]
-                    mask = curr_label != -100
-                    decoded_true = tokenizer.decode(curr_label[mask], skip_special_tokens=True)
-                    decoded_pred = tokenizer.decode(pred_tokens[row][mask], skip_special_tokens=True)
-
-                    accelerator.print(f"Predicted tokens: {decoded_pred}")
-                    accelerator.print(f"True tokens: {decoded_true}")
-
-
-            # gather loss before backprop in case of gradient accumulation
-            loss_values = accelerator.gather_for_metrics({"loss": loss.detach().float()})
-            train_loss.update(loss_values["loss"])
-
-            loss = loss / gradient_accumulation_steps
-            accelerator.backward(loss)
-            # get gradient norm of all params
-
-            # log LR in case something weird happens 
-            if config["wandb"]:
-                if step > 0 and step % (config["log_lr_every"] ) == 0:
-                    lr = optimizer.param_groups[0]["lr"]
-                    accelerator.log({"lr": lr}, step=curr_step)
-
-                    if learnable_alpha is False and isinstance(cross_attn_layer, int):
-                        curr_alpha = model.gpt_neox.layers[cross_attn_layer]._update_alpha(curr_step)
-                        accelerator.log({"alpha": curr_alpha}, step=curr_step)
-                    elif isinstance(cross_attn_layer, int):
-                        curr_alpha = model.gpt_neox.layers[cross_attn_layer].alpha.item()
-                        accelerator.log({"alpha": curr_alpha}, step=curr_step)
-
-            if (step + 1) % gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
-                optimizer.step()
-                if scheduler:
-                    scheduler.step()
-                optimizer.zero_grad()
-
-
-            if step > 0 and config["save_every"] > 0 and step % config["save_every"] == 0:
-                accelerator.save_state(f"{config['output_dir']}/step_{curr_step}")
-
-            if step > 0 and (step % config["eval_every"] == 0 or step == len(train_dataloader) - 1):
-                val_loss = evaluate(model, val_dataloader, step=curr_step, main_process=main_process)
-
-                log_train = {
-                        "train_loss": train_loss.compute()
-                    }
-                log_val = {
-                    "val_loss": val_loss.compute(),
-                }
-
-                if config["wandb"]:
-                    curr_step = step + epoch * len(train_dataloader)
-                    accelerator.log({**log_train, **log_val}, step=curr_step)
-
-                accelerator.print(f"Current LR: {optimizer.param_groups[0]['lr']}")
-                accelerator.print(format_metrics(log_train, "train", f" step {step} "))
-                accelerator.print(format_metrics(log_val, "val", f" step {step} "))
-
-                train_loss.reset()
-
-        accelerator.print(f"Epoch {epoch} finished")
-        accelerator.wait_for_everyone()
-        unwrapped_model = accelerator.unwrap_model(model)
-        if config["push_to_hub"]:
-            accelerator.print(f"Pushing to HF hub")
-            try:
-                if accelerator.is_main_process:
-                    unwrapped_model.push_to_hub(config["save_name"] + f"-epoch_{epoch}", private=True)
-
-            except Exception as e:
-                accelerator.print(e)
-                accelerator.print(f"Failed to push to hub")
-
-        unwrapped_model.save_pretrained(
-            f"{config['output_dir']}/epoch_{epoch}",
-            is_main_process=accelerator.is_main_process,
-            save_function=accelerator.save,
-            state_dict=accelerator.get_state_dict(model),
-        )
-            
-    accelerator.wait_for_everyone()
-    unwrapped_model = accelerator.unwrap_model(model)
-    unwrapped_model.save_pretrained(
-        f"{config['output_dir']}/final",
-        is_main_process=accelerator.is_main_process,
-        save_function=accelerator.save,
-        state_dict=accelerator.get_state_dict(model),
-    )
-
-    accelerator.end_training()
-
-    
-
-if __name__ == "__main__":
-    # parse arguments by reading in a config
-    parser = ArgumentParser()
-    parser.add_argument("--config", type=str, default="config.yaml")
-
-    args = parser.parse_args()
-
-    config = read_config(args.config)
-
-    if config["wandb"]:
-        accelerator = Accelerator(log_with="wandb")
-        accelerator.init_trackers(
-            project_name=config["wandb_project_name"],
-            config=config,
-            init_kwargs={"wandb": {"entity": config["wandb_entity"]}},
-        )
-    else:
-        accelerator = Accelerator()
-
-    train(accelerator, config=config)

gpt4all/utils/distributed_utils.py

@@ -1,27 +0,0 @@
-import torch.distributed as dist
-from contextlib import contextmanager
-
-
-def rank0_print(msg):
-    if dist.is_initialized():
-        if dist.get_rank() == 0:
-            print(msg)
-    else:
-        print(msg)
-
-
-@contextmanager
-def main_process_first(is_main):
-    yield from _goes_first(is_main)
-    
-
-        
-def _goes_first(is_main):
-    if not is_main:
-        dist.barrier()
-
-    yield
-
-    if is_main:
-        dist.barrier()
-    

head_node_setup.sh

@@ -0,0 +1,19 @@
+#!/bin/sh
+WORKER_IP=$1
+N_GPUS=$2
+
+sudo apt install -y nfs-kernel-server
+sudo mkdir -p ./data_multiplus
+sudo chmod 777 ./data_multiplus
+printf "${PWD}/data_multiplus ${WORKER_IP}(rw,sync,no_subtree_check)" | sudo tee -a /etc/exports
+sudo systemctl restart nfs-kernel-server
+
+sudo apt-get install -y pdsh
+export DSHPATH=$PATH
+export PDSH_RCMD_TYPE=ssh
+
+ssh-keygen -t rsa -N ''
+cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys
+
+sudo mkdir -p /job
+printf "localhost slots=$N_GPUS\n$WORKER_IP slots=$N_GPUS" | sudo tee /job/hostfile

inference.py

@@ -2,14 +2,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import torch.nn as nn
 from argparse import ArgumentParser
-from gpt4all.utils.read import read_config
-from accelerate.utils import set_seed
-from gpt4all.data.instruction_tuning_dataloader import load_data_for_inference
-from gpt4all.utils.distributed_utils import rank0_print
+from read import read_config
+from accelerate.utils import  set_seed
+from data import load_data_for_inference
 from tqdm import tqdm
-from datasets import Dataset
+from datasets import  Dataset
 import torch.distributed as dist
-from transformers.trainer_pt_utils import nested_numpify
+from transformers.trainer_pt_utils import  nested_numpify
 from transformers import DefaultDataCollator
 from torch.utils.data import DataLoader, DistributedSampler
 import numpy as np
@@ -22,64 +21,56 @@ def calc_cross_entropy_no_reduction(lm_logits, labels):
     shift_logits = lm_logits[..., :-1, :].contiguous()
     shift_labels = labels[..., 1:].contiguous()
     # Flatten the tokens
-    loss_fct = nn.CrossEntropyLoss(reduction="none")
+    loss_fct = nn.CrossEntropyLoss(reduction='none')
     loss = loss_fct(shift_logits.permute(0, 2, 1), shift_labels).mean(dim=1)
 
     return loss
 
 
+def rank0_print(msg):
+    if dist.get_rank() == 0:
+        print(msg)
+        
+
 def inference(config):
-    set_seed(config["seed"])
+    set_seed(config['seed'])
 
     rank0_print(f"World size: {dist.get_world_size()}")
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        config["tokenizer_name"], model_max_length=config["max_length"]
-    )
+    tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'], model_max_length=config['max_length'])
     # llama has no pad token, set it to new token
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    train_dataset, val_dataset = load_data_for_inference(config, tokenizer)
+        
+    train_dataset, val_dataset = load_data_for_inference(config, tokenizer) 
 
     num_processes = dist.get_world_size()
     local_rank = dist.get_rank()
 
-    train_sampler = DistributedSampler(
-        train_dataset,
-        shuffle=False,
-        drop_last=True,
-        num_replicas=num_processes,
-        rank=local_rank,
-    )
+    train_sampler = DistributedSampler(train_dataset, shuffle=False, drop_last=True, num_replicas=num_processes, rank=local_rank)
     train_dataloader = DataLoader(
         train_dataset,
         collate_fn=DefaultDataCollator(),
         batch_size=config["batch_size"],
         sampler=train_sampler,
-        drop_last=True,
+        drop_last=True
     )
 
-    val_sampler = DistributedSampler(
-        val_dataset,
-        shuffle=False,
-        drop_last=True,
-        num_replicas=num_processes,
-        rank=local_rank,
-    )
+    val_sampler = DistributedSampler(val_dataset, shuffle=False, drop_last=True, num_replicas=num_processes, rank=local_rank)
     val_dataloader = DataLoader(
         val_dataset,
         collate_fn=DefaultDataCollator(),
         batch_size=config["batch_size"],
         sampler=val_sampler,
-        drop_last=True,
+        drop_last=True
     )
 
-    model = AutoModelForCausalLM.from_pretrained(
-        config["model_name"],
-        trust_remote_code=True,
-        torch_dtype=torch.bfloat16,
-    )
+
+    model = AutoModelForCausalLM.from_pretrained(config["model_name"], 
+                                                    trust_remote_code=True,
+                                                    torch_dtype=torch.bfloat16,
+                                                    ) 
     model.to(f"cuda:{local_rank}")
 
     with torch.no_grad():
@@ -87,18 +78,14 @@ def inference(config):
         for batch in tqdm(train_dataloader, disable=local_rank != 0):
             batch["input_ids"] = batch["input_ids"].to(f"cuda:{local_rank}")
             batch["labels"] = batch["labels"].to(f"cuda:{local_rank}")
-            outputs = model(
-                input_ids=batch["input_ids"],
-                labels=batch["labels"],
-                output_hidden_states=True,
-            )
+            outputs = model(input_ids=batch["input_ids"], labels=batch["labels"], output_hidden_states=True)
             loss = calc_cross_entropy_no_reduction(outputs.logits, batch["labels"])
             train_outputs["loss"].extend(loss)
 
             embeddings = outputs.hidden_states[-1]
             batch_size = batch["input_ids"].shape[0]
             sequence_lengths = []
-            # since we use mutiturn with multiple <|endoftext|>, we need to find the place where
+            # since we use mutiturn with multiple <|endoftext|>, we need to find the place where 
             # <|endoftext|> is repeated
             for item in batch["input_ids"]:
                 indices = torch.where(item == tokenizer.pad_token_id)[0]
@@ -114,9 +101,7 @@ def inference(config):
                     sequence_lengths.append(len(item) - 1)
 
             sequence_lengths = torch.tensor(sequence_lengths)
-            pooled_logits = embeddings[
-                torch.arange(batch_size, device=embeddings.device), sequence_lengths
-            ]
+            pooled_logits = embeddings[torch.arange(batch_size, device=embeddings.device), sequence_lengths]
 
             train_outputs["embeddings"].append(pooled_logits)
             train_outputs["index"].extend(batch["index"].to(model.device))
@@ -130,63 +115,51 @@ def inference(config):
         train_outputs["embeddings"] = np.concatenate(train_outputs["embeddings"])
 
         df_train = Dataset.from_dict(train_outputs)
+        df_train = df_train.sort("index")
         curr_idx = df_train["index"]
 
         # compute mask in pyarrow since it's super fast
         # ty @bmschmidt for showing me this!
         table = train_dataset.data
-        mask = pc.is_in(table["index"], value_set=pa.array(curr_idx, pa.int32()))
+        mask = pc.is_in(table['index'], value_set=pa.array(curr_idx, pa.int32()))
         filtered_table = table.filter(mask)
         # convert from pyarrow to Dataset
         filtered_train = Dataset.from_dict(filtered_table.to_pydict())
 
         filtered_train = filtered_train.add_column("embeddings", df_train["embeddings"])
         filtered_train = filtered_train.add_column("loss", df_train["loss"])
-        filtered_train = filtered_train.add_column(
-            "is_train", [True] * len(filtered_train)
-        )
+        filtered_train = filtered_train.add_column("is_train", [True] * len(filtered_train))
 
-        filtered_train.to_json(
-            f"inference/epoch_2_embeddings_train_shard_{local_rank}.jsonl",
-            lines=True,
-            orient="records",
-            num_proc=64,
-        )
+        filtered_train.to_json(f"inference/epoch_2_embeddings_train_shard_{local_rank}.jsonl", lines=True, orient="records", num_proc=64)
 
         val_outputs = {"loss": [], "embeddings": [], "index": []}
         for batch in tqdm(val_dataloader, disable=local_rank != 0):
             batch["input_ids"] = batch["input_ids"].to(f"cuda:{local_rank}")
             batch["labels"] = batch["labels"].to(f"cuda:{local_rank}")
-            outputs = model(
-                input_ids=batch["input_ids"],
-                labels=batch["labels"],
-                output_hidden_states=True,
-            )
+            outputs = model(input_ids=batch["input_ids"], labels=batch["labels"])
             loss = calc_cross_entropy_no_reduction(outputs.logits, batch["labels"])
             val_outputs["loss"].extend(loss)
 
-            embeddings = outputs.hidden_states[-1]
+            logits = outputs.logits
             batch_size = batch["input_ids"].shape[0]
             sequence_lengths = []
-            # since we use mutiturn with multiple <|endoftext|>, we need to find the place where
+            # since we use mutiturn with multiple <|endoftext|>, we need to find the place where 
             # <|endoftext|> is repeated
             for item in batch["input_ids"]:
                 indices = torch.where(item == tokenizer.pad_token_id)[0]
                 found = False
                 for index in indices:
-                    # case where sequence is less than max length
                     if torch.all(item[index:] == tokenizer.pad_token_id):
                         sequence_lengths.append(index)
                         found = True
                         break
-                # case where sequence is >= max length
+
+                # no match found
                 if not found:
                     sequence_lengths.append(len(item) - 1)
 
             sequence_lengths = torch.tensor(sequence_lengths)
-            pooled_logits = embeddings[
-                torch.arange(batch_size, device=embeddings.device), sequence_lengths
-            ]
+            pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
 
             val_outputs["embeddings"].append(pooled_logits)
             val_outputs["index"].extend(batch["index"].to(model.device))
@@ -199,26 +172,23 @@ def inference(config):
         val_outputs["embeddings"] = np.concatenate(val_outputs["embeddings"])
 
         df_val = Dataset.from_dict(val_outputs)
+        df_val = df_val.sort("index")
         curr_idx = df_val["index"]
 
         # compute mask in pyarrow since it's super fast
         # ty @bmschmidt for showing me this!
         table = val_dataset.data
-        mask = pc.is_in(table["index"], value_set=pa.array(curr_idx, pa.int32()))
+        mask = pc.is_in(table['index'], value_set=pa.array(curr_idx, pa.int32()))
         filtered_table = table.filter(mask)
         # convert from pyarrow to Dataset
         filtered_val = Dataset.from_dict(filtered_table.to_pydict())
+
         filtered_val = filtered_val.add_column("embeddings", df_val["embeddings"])
         filtered_val = filtered_val.add_column("loss", df_val["loss"])
         filtered_val = filtered_val.add_column("is_train", [False] * len(filtered_val))
 
-        filtered_val.to_json(
-            f"inference/epoch_2_embeddings_val_shard_{local_rank}.jsonl",
-            lines=True,
-            orient="records",
-            num_proc=64,
-        )
-
+        filtered_val.to_json(f"inference/epoch_2_embeddings_val_shard_{local_rank}.jsonl", lines=True, orient="records", num_proc=64)
+    
 
 def main():
     dist.init_process_group("nccl")
@@ -234,3 +204,4 @@ def main():
 if __name__ == "__main__":
     # parse arguments by reading in a config
     main()
+    

launcher.sh

@@ -1,88 +0,0 @@
-#!/bin/bash
-
-# Display header
-echo "=========================================================="
-echo " ██████  ██████  ████████ ██   ██  █████  ██      ██      "
-echo "██       ██   ██    ██    ██   ██ ██   ██ ██      ██      "
-echo "██   ███ ██████     ██    ███████ ███████ ██      ██      "
-echo "██    ██ ██         ██         ██ ██   ██ ██      ██      "
-echo " ██████  ██         ██         ██ ██   ██ ███████ ███████ "
-echo " └─> https://github.com/nomic-ai/gpt4all"
-
-# Function to detect macOS architecture and set the binary filename
-detect_mac_arch() {
-  local mac_arch
-  mac_arch=$(uname -m)
-  case "$mac_arch" in
-    arm64)
-      os_type="M1 Mac/OSX"
-      binary_filename="gpt4all-lora-quantized-OSX-m1"
-      ;;
-    x86_64)
-      os_type="Intel Mac/OSX"
-      binary_filename="gpt4all-lora-quantized-OSX-intel"
-      ;;
-    *)
-      echo "Unknown macOS architecture"
-      exit 1
-      ;;
-  esac
-}
-
-# Detect operating system and set the binary filename
-case "$(uname -s)" in
-  Darwin*)
-    detect_mac_arch
-    ;;
-  Linux*)
-    if grep -q Microsoft /proc/version; then
-      os_type="Windows (WSL)"
-      binary_filename="gpt4all-lora-quantized-win64.exe"
-    else
-      os_type="Linux"
-      binary_filename="gpt4all-lora-quantized-linux-x86"
-    fi
-    ;;
-  CYGWIN*|MINGW32*|MSYS*|MINGW*)
-    os_type="Windows (Cygwin/MSYS/MINGW)"
-    binary_filename="gpt4all-lora-quantized-win64.exe"
-    ;;
-  *)
-    echo "Unknown operating system"
-    exit 1
-    ;;
-esac
-echo "================================"
-echo "== You are using $os_type."
-
-
-# Change to the chat directory
-cd chat
-
-# List .bin files and prompt user to select one
-bin_files=(*.bin)
-echo "== Available .bin files:"
-for i in "${!bin_files[@]}"; do
-  echo "   [$((i+1))] ${bin_files[i]}"
-done
-
-# Function to get user input and validate it
-get_valid_user_input() {
-  local input_valid=false
-
-  while ! $input_valid; do
-    echo "==> Please enter a number:"
-    read -r user_selection
-    if [[ $user_selection =~ ^[0-9]+$ ]] && (( user_selection >= 1 && user_selection <= ${#bin_files[@]} )); then
-      input_valid=true
-    else
-      echo "Invalid input. Please enter a number between 1 and ${#bin_files[@]}."
-    fi
-  done
-}
-
-get_valid_user_input
-selected_bin_file="${bin_files[$((user_selection-1))]}"
-
-# Run the selected .bin file with the appropriate command
-./"$binary_filename" -m "$selected_bin_file"

requirements.txt

@@ -2,7 +2,7 @@ accelerate
 datasets
 torchmetrics
 evaluate
-transformers>=4.28.0
+transformers
 wandb
 pip
 peft
@@ -10,9 +10,4 @@ nodelist-inflator
 deepspeed
 sentencepiece
 jsonlines
-nomic
-scikit-learn
-matplotlib
-apache_beam 
-mwparserfromhell
-hnswlib
+nomic

setup.py

@@ -1,34 +0,0 @@
-from setuptools import setup, find_packages
-
-with open('README.md', 'r', encoding='utf-8') as f:
-    long_description = f.read()
-
-with open('requirements.txt', 'r', encoding='utf-8') as f:
-    requirements = [line.strip() for line in f if line.strip()]
-
-setup(
-    name='gpt4all',
-    version='0.0.1',
-    author='nomic-ai',
-    author_email='zach@nomic-ai',
-    description='an ecosystem of open-source chatbots trained on a massive collections of clean assistant data including code, stories and dialogue',
-    long_description=long_description,
-    long_description_content_type='text/markdown',
-    url='https://github.com/nomic-ai/gpt4all',
-    packages=find_packages(),
-    install_requires=requirements,
-    classifiers=[
-        'Development Status :: 3 - Alpha',
-        'License :: OSI Approved :: MIT License',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.6',
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
-        'Programming Language :: Python :: 3.9',
-        'Topic :: Text Processing :: Linguistic',
-        'Topic :: Scientific/Engineering :: Artificial Intelligence',
-        'Intended Audience :: Science/Research',
-        'Operating System :: OS Independent',
-    ],
-    python_requires='>=3.6',
-)

train.py

@@ -1,13 +1,13 @@
 import os
-from transformers import AutoModelForCausalLM, AutoTokenizer, get_scheduler
+from transformers import AutoModelForCausalLM, AutoTokenizer, get_scheduler, LlamaForCausalLM
 import torch
 from torch.optim import AdamW
 from argparse import ArgumentParser
-from gpt4all.utils.read import read_config
+from read import read_config
 from accelerate import Accelerator
 from accelerate.utils import DummyScheduler, DummyOptim, set_seed
 from peft import get_peft_model, LoraConfig, TaskType
-from gpt4all.data.instruction_tuning_dataloader import load_data
+from data import load_data
 from torchmetrics import MeanMetric
 from tqdm import tqdm
 import wandb
@@ -100,11 +100,11 @@ def train(accelerator, config):
             name="cosine",
             optimizer=optimizer,
             num_warmup_steps=config["warmup_steps"] * accelerator.num_processes,
-            num_training_steps=total_num_steps,
+            num_training_steps=total_num_steps * accelerator.num_processes,
         )
     else:
         scheduler = DummyScheduler(
-            optimizer, total_num_steps=total_num_steps, warmup_num_steps=config["warmup_steps"]
+            optimizer, total_num_steps=config["warmup_steps"], warmup_num_steps=config["warmup_steps"]
         )
 
     model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
@@ -192,7 +192,7 @@ def train(accelerator, config):
             accelerator.print(f"Failed to push to hub")
 
         unwrapped_model.save_pretrained(
-            f"{config['output_dir']}/epoch_{epoch}",
+            f"{config['output_dir']}/-epoch_{epoch}",
             is_main_process=accelerator.is_main_process,
             save_function=accelerator.save,
             state_dict=accelerator.get_state_dict(model),

worker_node_setup.sh

@@ -0,0 +1,6 @@
+#!/bin/sh
+HEAD_IP=$1
+
+sudo apt install -y nfs-common
+sudo mkdir -p ./data_multiplus
+sudo mount ${HEAD_IP}:${PWD}/data_multiplus ./data_multiplus