LongMemEval

repository·main·Indexed 19 days ago

https://github.com/xiaowu0162/longmemeval

A benchmark for testing the long-term memory capabilities of chat assistants across five core abilities: Information Extraction, Multi-Session Reasoning, Knowledge Updates, Temporal Reasoning, and Abstention. It includes datasets in small and medium scales, tools for memory indexing and retrieval (BM25, Contriever, Stella, GTE), and an automated QA evaluation pipeline using metric models like GPT-4o and Llama-3.1-70B-Instruct.

Tokens
12.1K
Snippets
33
Records
42
Agent score
75%

What's inside LongMemEval

  1. Understand the LongMemEval dataset format

    main

    The dataset contains three main files:

    • longmemeval_s.json: Small scale (approx. 40 history sessions, ~115k tokens).
    • longmemeval_m.json: Medium scale (approx. 500 sessions).
    • longmemeval_oracle.json: Oracle retrieval version (only contains evidence sessions).

    Each instance in the JSON files contains:

    • question_id: Unique identifier.
    • question_type: One of single-session-user, single-session-assistant, single-session-preference, temporal-reasoning, knowledge-update, and multi-session. (Note: if question_id ends in _abs, it is an abstention type).
    • question: The question text.
    • answer: The ground truth answer.
    • question_date: Date of the question.
    • haystack_session_ids: List of session IDs in the history.
    • haystack_dates: Timestamps for the history sessions.
    • haystack_sessions: The chat history. A list of sessions, where each session is a list of turns: {"role": "user"|"assistant", "content": "..."}. Turns with required evidence include "has_answer": true.
    • answer_session_ids: List of session IDs containing the evidence.
  2. Download LongMemEval dataset

    main

    The dataset is available on HuggingFace. Download the files and uncompress them into a data/ directory at the root of your project.

    mkdir -p data/
    cd data/
    wget https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_oracle.json
    wget https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json
    wget https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_m_cleaned.json
    cd ..
  3. Run Baseline Memory Retrieval

    main

    To perform memory indexing and retrieval, use run_retrieval.sh in the src/retrieval directory.

    Arguments:

    • IN_FILE: Path to the input file.
    • RETRIEVER: One of flat-bm25, flat-contriever, flat-stella (requires manual download of Stella V5 1.5B), or flat-gte (gte-Qwen2-7B-instruct).
    • GRANULARITY: turn or session.

    Notes:

    • Dense embedding models utilize all available GPUs by default.
    • Results are output to retrieval_logs/.
    • To view metrics from a log file: python3 src/evaluation/print_retrieval_metrics.py log_file.
    • Evaluation skips the 30 abstention instances (non-existing events).
    cd src/retrieval
    bash run_retrieval.sh IN_FILE RETRIEVER GRANULARITY
  4. Run Long-Context Generation experiments

    main

    To run the long-context generation baseline (where the model is provided with the full history), use run_generation.sh in the src/generation directory.

    Arguments:

    • DATA_FILE: Path to a released JSON file. Note: longmemeval_s.json and longmemeval_oracle.json fit in 128k context; longmemeval_m.json is too long for this method.
    • MODEL: Alias of the model (configured in run_generation.sh).
    • full-history-session: The session type.
    • TOPK: Max number of history sessions to provide (recommend 1000).
    • HISTORY_FORMAT: json (recommended) or nl.
    • USERONLY: true or false (to remove assistant messages; recommend false).
    • READING_METHOD: direct, con, or con-separate (recommend con for extraction then reasoning).

    Logs are saved to generation_logs/.

    cd src/generation
    bash run_generation.sh DATA_FILE MODEL full-history-session TOPK [HISTORY_FORMAT] [USERONLY] [READING_METHOD]
  5. Run Retrieval-Augmented Generation (RAG)

    main

    To perform question answering using retrieved memory, use run_generation.sh in the src/generation directory.

    Arguments:

    • RETRIEVAL_LOG_FILE: The output file from the retrieval step (must contain the retrieval_results field).
    • MODEL: Alias of the model.
    • EXP: The expansion identifier in the format [RETRIEVER]-[GRANULARITY] (e.g., flat-stella-session).
    • TOPK: Max number of history sessions.
    • HISTORY_FORMAT: json or nl.
    • USERONLY: true or false.
    • READING_METHOD: direct, con, or con-separate.
    cd src/generation
    bash run_generation.sh RETRIEVAL_LOG_FILE MODEL EXP TOPK [HISTORY_FORMAT] [USERONLY] [READING_METHOD]
  6. Perform Time-Aware Query Expansion

    main

    This process prunes the search space by using timestamped events and inferred time ranges.

    1. Download the extracted timestamped events and unzip them into LongMemEval/index_expansion_logs/.
    2. Run the pruning script:

    Arguments:

    • TIMESTAMP_EVENT_FILE: The downloaded timestamped events file.
    • RETRIEVAL_LOG: Output from a previous retrieval experiment.
    • GRANULARITY: session or turn (must match the other two arguments).
    cd src/index_expansion
    python3 temp_query_search_pruning.py TIMESTAMP_EVENT_FILE RETRIEVAL_LOG GRANULARITY
  7. Run Retrieval with Index Expansion

    main

    To run experiments using key expansion, first download the expansion outputs and place them in LongMemEval/index_expansion_logs/. Then run run_retrieval.sh with additional arguments.

    Additional Arguments:

    • EXPANSION_TYPE: session-summ, session-keyphrase, session-userfact, turn-keyphrase, or turn-userfact.
    • JOIN_MODE:
      • separate: Add a new (key, value) pair.
      • merge: Merge expansion with original key.
      • replace: Replace original key with expansion.
    • CACHE: Path to the cache file corresponding to the EXPANSION_TYPE.
    cd src/retrieval
    bash run_retrieval.sh IN_FILE RETRIEVER GRANULARITY EXPANSION_TYPE JOIN_MODE CACHE
  8. Serve an open-weight LLM via vLLM

    main

    To use an open-weight reader LLM, you can serve it locally using an OpenAI API emulator via vllm. This is required if you are not using OpenAI's direct API.

    Use serve_vllm.sh for standard serving or serve_vllm_with_maxlen.sh if you need to limit the maximum number of tokens due to memory constraints.

    # Standard serving
    cd src/utils
    bash serve_vllm.sh GPU MODEL PORT TP_SIZE
    
    # Serving with max token limit
    bash serve_vllm_with_maxlen.sh GPU MODEL MAXLEN PORT TP_SIZE
  9. Reproduce LongMemEval history compilation

    main

    You can use sample_haystack_and_timestamp.py to reproduce the history compilation process.

    Command Syntax: python sample_haystack_and_timestamp.py task n_questions min_n_haystack_filler max_n_haystack_filler enforce_json_length

    Task Name Mapping:

    • single_hop $\rightarrow$ single-session-user
    • implicit_preference_v2 $\rightarrow$ single-session-preference
    • assistant_previnfo $\rightarrow$ single-session-assistant
    • two_hop $\rightarrow$ multi-session
    • multi_session_synthesis $\rightarrow$ multi-session
    • temp_reasoning_implicit $\rightarrow$ temporal-reasoning
    • temp_reasoning_explicit $\rightarrow$ temporal-reasoning
    • knowledge_update $\rightarrow$ knowledge-update

    Parameter Guide:

    • n_questions: Max number of questions.
    • min_n_haystack_filler / max_n_haystack_filler: Limits on session counts (e.g., 80 for longmemeval_s, 500 for longmemeval_m).
    • enforce_json_length: Limits chat history length (e.g., 115000 for longmemeval_s).
  10. Install LongMemEval for evaluation only

    main

    If you only need to calculate metrics on outputs produced by your own system using src/evaluation/evaluate_qa.py, use the minimal requirement set. This is the lightweight option for users who do not need to run the specific memory systems described in the paper.

    conda create -n longmemeval-lite python=3.9
    conda activate longmemeval-lite
    pip install -r requirements-lite.txt
  11. Test your system on LongMemEval

    main

    To evaluate your system:

    1. Feed the timestamped history to your chat system.
    2. Save your outputs in a .jsonl format where each line contains question_id and hypothesis.
    3. Run the evaluation script.

    Note: You must set OPENAI_API_KEY and optionally OPENAI_ORGANIZATION environment variables.

    export OPENAI_API_KEY=YOUR_API_KEY
    export OPENAI_ORGANIZATION=YOUR_ORGANIZATION
    cd src/evaluation
    python3 evaluate_qa.py gpt-4o your_hypothesis_file ../../data/longmemeval_oracle.json
  12. Aggregate evaluation metrics from logs

    main

    The evaluate_qa.py script generates a log file named [your_hypothesis_file].log. To aggregate the scores from this log, use print_qa_metrics.py.

    # Assuming you are in the src/evaluation folder
    python3 print_qa_metrics.py gpt-4o your_hypothesis_file.log ../../data/longmemeval_oracle.json