feat: Add nested KV and doc QA experiment scripts for paper

paper_experiments/README.md

@@ -0,0 +1,47 @@
+
+## Nested K/V (`nested_kv_task`)
+This task runs K/V lookups on synthetic data. You can run it with `icml_experiments/nested_kv_task/run.sh`.
+
+## Document Q/A (`doc_qa_task`)
+This task runs question answering on a set of embedded wikipedia passages.
+
+### Setup
+You need a a running postgres database to run this experiment and an OpenAI account. Set your enviornment variables:
+```
+export PGVECTOR_TEST_DB_URL=postgresql+pg8000://{username}:{password}@localhost:8888/{db}
+export OPENAI_API_KEY={key}
+```
+
+## Download data
+Download the wikipedia embedding at:
+```
+huggingface-cli download nlpkevinl/wikipedia_openai_embeddings --repo-type dataset
+```
+
+## Loading embeddings
+Run the script `./0_load_embeddings.sh`.
+
+This step will take a while. You can check the status of the loading by connecting to `psql`:
+```
+> psql -h localhost -p {password} -U {username} -d {db}
+> SELECT COUNT(*) FROM memgpt_passages;
+```
+Once completed, there will be ~19 million rows in the database.
+
+### Creating an index
+To avoid extremeley slow queries, you need to create an index:
+```
+CREATE INDEX ON memgpt_passages USING hnsw (embedding vector_l2_ops);
+```
+You can check to see if the index was created successfully with:
+```
+> SELECT indexname, indexdef FROM pg_indexes WHERE tablename = 'memgpt_passages';
+
+memgpt_passages_embedding_idx | CREATE INDEX memgpt_passages_embedding_idx ON public.memgpt_passages USING hnsw (embedding vector_cosine_ops) WITH (m='24', ef_construction='100')
+```
+
+## Running Document Q/A
+Run the script `./1_run_docqa.sh {model_name} {n_docs} {memgpt/model_name}`.
+
+## Evaluation
+Run the script `./2_run_eval.sh`.

paper_experiments/doc_qa_task/0_load_embeddings.sh

@@ -0,0 +1,17 @@
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-06.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-07.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-08.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-09.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-01.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-02.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-03.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-04.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-05.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-06.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-07.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_2-08.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-01.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-02.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-03.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-04.jsonl
+python load_wikipedia_embeddings.py --file data/wikipedia_passages_shard_1-05.jsonl

paper_experiments/doc_qa_task/1_run_docqa.sh

@@ -0,0 +1,4 @@
+docs=$2
+model=$1
+baseline=$3
+python icml_experiments/doc_qa_task/doc_qa.py --model $model --baseline $baseline --num_docs $docs

paper_experiments/doc_qa_task/2_run_eval.sh

@@ -0,0 +1,18 @@
+docs=(1 5 10 20 50 100 200 700)
+models=("gpt-4-0613" "gpt-3.5-turbo-1106" "gpt-4-1106-preview")
+
+## run memgpt eval
+for model in  "${models[@]}";
+do
+  poetry run python icml_experiments/doc_qa_task/llm_judge_doc_qa.py --file results/doc_qa_results_model_${model}.json
+done
+
+# Iterate over each model
+for model in "${models[@]}"; do
+    # Iterate over each doc
+    for doc in "${docs[@]}"; do
+        # Construct and run the command
+        echo "Running for model $model with $doc docs..."
+        poetry run python icml_experiments/doc_qa_task/llm_judge_doc_qa.py --file results/doc_qa_baseline_model_${model}_num_docs_${doc}.json --baseline
+    done
+done