unlimiformer: TypeError: torch_replacement_knn_gpu() got an unexpected keyword argument 'device'

Just for future reference, command line did you use to generate the last output? (I am curious about the exact model, exact layer_begin, and exact prompt)

Best, Uri

here’s the command for the 80k tokens input:

python src/run_generation.py --model_type llama --model_name_or_path meta-llama/Llama-2-7b-chat-hf \
    --prefix "<<SYS>>\n You are a helpful assistant. Answer with detailed responses according to the entire instruction or question. \n<</SYS>>\n\n [INST] Summarize the following book: " \
    --prompt example_inputs/1.txt \
    --suffix " [/INST]" --test_unlimiformer --fp16 --length 200 --layer_begin 20 \
    --index_devices 1 --datastore_device 2 --stream_output

And surprisingly, I modified the instruction a little in the prefix and using a 135k tokens input to test again. Here’s the result:

=== GENERATED SEQUENCE 1 (input length: 135830) ===
||| This is a cookbook with recipes for various types of salads. The book is divided into seasons, with a chapter for each season, and each chapter contains a variety of salad recipes, each with a brief description of the dish and an explanation of how to prepare it. Each recipe is accompanied by notes and suggestions for complementary dishes, and wine or beer pairings.

* Winter: Three-Alarm Salad, Galette of Greens and Goat cheese, with Braised Mushrooms and Mustard, Bitter greens, carrot and Pickled Zucchini.
* Spring: Sautéed Chicken, Kitchen Garden Greens, Chèvre, and Chive. Spring Onion Tarts, Toppings, 45
* Summer: Zucchettis, Cauliflower, Ratatelli, Chiveroli, and Chilled, and Lemon Dress.

Amazing! And the command for above:

python src/run_generation.py --model_type llama --model_name_or_path meta-llama/Llama-2-7b-chat-hf \
    --prefix "<<SYS>>\n You are a helpful assistant. Answer with detailed responses according to the entire instruction or question. \n<</SYS>>\n\n [INST] Summarize the content of the following book: " \
    --prompt example_inputs/1.txt \
    --suffix " [/INST]" --test_unlimiformer --fp16 --length 200 --layer_begin 20 \
    --index_devices 1 --datastore_device 2 --stream_output

But it’s weird that the model’s output for summarizing harry potter is still strange although using the same flags as the above cookbook:

=== GENERATED SEQUENCE 1 (input length: 131318) ===
||| Harry Potter and the Philosopher's Stone, J.K. Rowling.


He had been famous -- Harry -- since he'd become the
magic's hero... He'd been famous in front of the Muggles, too. 
It was a weird feeling, famous -- being a footloos an' talked to'...' 



They'd worked out how to get past Fluffy without trying

"...Midgit... it was Death E ly... Theoretical..."

(Yes, the model indeed output a lot of \n, maybe it’s harry’s magic I guess…)

Ok finally got it to work in Google Colab on an A100 40G. For anyone curious I used StableBeluga-13B and it took around 9 minutes to get a summary of Harry Potter, which is pretty good especially since you can’t even fit the full book in Claude 100k! I’m thoroughly impressed!

Here is the Code I used to get it working in Colab:

First, in order to get the latest version of Faiss you have to upgrade Python to 3.10 since it’s automatically set to 3.7

 !wget https://github.com/korakot/kora/releases/download/v0.10/py310.sh
!bash ./py310.sh -b -f -p /usr/local
!python -m ipykernel install --name "py310" --user

Then you’ll want to install mini-conda so that you can use install faiss using conda.

################################################################################
# INSTALL CONDA ON GOOGLE COLAB
################################################################################
! wget https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-Linux-x86_64.sh
! chmod +x Miniconda3-py310_23.3.1-0-Linux-x86_64.sh
! bash ./Miniconda3-py310_23.3.1-0-Linux-x86_64.sh -b -f -p /usr/local
import sys
sys.path.append('/usr/local/lib/python3.10/site-packages/')

Install using conda-forge

!conda install -c conda-forge faiss-gpu -y

Then clone the repo in colab

!git clone https://github.com/abertsch72/unlimiformer.git

Install the requirements. (In this instance some of these aren’t required but I liked to have them just in case)

%pip install -r requirements.txt
%pip install -q -U bitsandbytes
%pip install -q -U git+https://github.com/huggingface/transformers.git
%pip install -q -U git+https://github.com/huggingface/peft.git
%pip install -q -U git+https://github.com/huggingface/accelerate.git
%pip install -q datasets
%pip install tensorrt

Cd into the src folder in Unlimiformer

%cd /content/unlimiformer/src

Then you should be good to run the script! Just be sure that the --index_devices and --datastore_device are set correctly. In my case I set them to 0

!python run_generation.py --model_type llama --model_name_or_path stabilityai/StableBeluga-13B \
    --prefix "### System:\n You are a helpful assistant. Answer with detailed responses according to the entire instruction or question. \n### User:\n\n [INST] Summarize the following book: " \
    --prompt example_inputs/harry_potter_full.txt \
    --suffix " ### Assistant" --test_unlimiformer --fp16 --length 200 --layer_begin 22 \
    --index_devices 0 --datastore_device 0

This worked pretty well after I set the --layer_begin to 22 (a little over half the number of layers in the model). Here’s the summary:

Harry Potter is a young boy who discovers he is a wizard, invited to attend the Hogwarts School of Witchcraft and Wizardry. He embarks on an adventure with his friends Ronald Weasley and Hermione Granger to face various challenges and enemies such as Voldemort and Lord Voldemort's supporters. Their journey involves discovering their true identities, unraveling mysteries, and learning valuable lessons about friendship, courage, and the fight against evil.</s>

Thanks again for all the hard work you and your team did @urialon I’m pretty hyped about this!