DOCSEARCHbyAlgolia
Submit Docs

AriGraph

repository·main·Indexed 41 days ago

https://github.com/airi-institute/arigraph

An external memory architecture for LLM agents utilizing a semantic knowledge graph augmented with episodic memory. Optimized for RAG performance in multi-hop question-answering and text-based games, featuring implementations for Contriever and XLMRetriever with configurable pooling strategies.

Tokens
1.2K
Snippets
3
Records
5
Agent score
2%

What's inside AriGraph

  1. Explore the AriGraph repository structure

    main

    The repository is organized into the following functional directories:

    • agents/: Contains the GPTagent implementation.
    • envs/: Contains TextWorld files used for environment loading.
    • qa_data/: Contains QA datasets (e.g., MuSiQue, HotpotQA) used for testing.
    • graphs/: Contains the TripletGraph (defined in parent_graph.py) and other graph implementations that inherit from it.
    • logs/: Stores logs from every reported agent run.
    • prompts/: Contains the prompts used within the various pipelines.
    • src/ and utils/: Contain service classes and utility functions.

    Note: Individual Python files in the root or other directories typically contain execution pipelines for agents or code for running games in interactive console mode. These pipelines include highlighted, changeable sections for execution parameters.

  2. Understand the AriGraph architecture

    main

    AriGraph serves as an external memory architecture for Large Language Models (LLMs). It functions by constructing a semantic knowledge graph that incorporates episodic vertices and edges. This structure is designed to improve Retrieval-Augmented Generation (RAG) performance, particularly in complex environments like text-based games.

    Key components include:

    • Semantic Knowledge Graph: The core memory structure.
    • Episodic Memory: Added vertices and edges that capture temporal/event-based information to enhance the world model.

  3. Install AriGraph requirements

    main

    AriGraph requires a Linux system due to TextWorld dependencies. On Debian/Ubuntu-based systems, you must install specific system libraries before installing Python dependencies.

    System Dependencies

    Debian/Ubuntu:

    sudo apt update && sudo apt install build-essential libffi-dev python3-dev curl git

    macOS:

    brew install libffi curl git

    Python Dependencies

    Ensure you are using Python 3.11+, then install the project requirements:

    pip install -r requirements.txt
    sudo apt update && sudo apt install build-essential libffi-dev python3-dev curl git

pip install -r requirements.txt

  4. Load a retriever model with load_retriever()

    main

    Use load_retriever(model_path, pooling="average", random_init=False) to initialize a semantic retriever model and its corresponding tokenizer.

    This function supports two modes:

    1. Checkpoint Loading: If model_path/checkpoint.pth exists, it loads the weights and metadata from the checkpoint. It automatically detects whether to use the Contriever or XLMRetriever class based on the retriever_model_id found in the checkpoint or the presence of 'xlm' in the ID.
    2. HuggingFace/Local Path Loading: If no checkpoint is found, it treats model_path as a HuggingFace model ID or a local directory containing a transformer model.

    Returns a tuple of (retriever, tokenizer, retriever_model_id).

    from src.contriever import load_retriever

# Load from a checkpoint directory
retriever, tokenizer, model_id = load_retriever("/path/to/checkpoint")

# Or load from a HuggingFace model ID
retriever, tokenizer, model_id = load_retriever("bert-base-uncased")

  5. Configure pooling in Contriever and XLMRetriever

    main

    Both Contriever and XLMRetriever support different pooling strategies to convert token-level hidden states into a single embedding vector. This is configured via the pooling argument during initialization or via the model's config.pooling attribute.

    Supported pooling options:

    • "average": Computes the mean of the last hidden states across the sequence length (masked by attention_mask).
    • "cls": Uses the first token ([CLS]) from the last hidden state.

    Additionally, the forward method accepts a normalize boolean flag. If set to True, the resulting embedding is normalized using torch.nn.functional.normalize along the last dimension.

    # Example of how pooling and normalization are applied during forward pass
# (Conceptual usage based on class definitions)

# Using average pooling (default)
embeddings = retriever(input_ids, attention_mask=attention_mask, normalize=True)

# The pooling type is determined by the config set during initialization
# retriever = Contriever(config, pooling="cls")