rea

Environment Setup

To ensure reproducibility, please use the following environment configuration:

• transformers >= 4.36.0
• torch >= 2.1.0
• accelerate
• peft (for dual-LoRA loading)

Usage Example

The following code demonstrates how to use REA to generate responses with integrated retrieval and memory management.

from transformers import AutoModel, GenerationConfig
import torch

# Path to the anonymous model repository
llm_model_path = "acl2026/rea"

# Load model with custom retrieval logic
model = AutoModel.from_pretrained(
    llm_model_path,
    local_files_only=True,
    trust_remote_code=True,  # Required for custom memory management and HCR logic
    torch_dtype=torch.bfloat16,
).to('cuda')

# Conversation history (document) and current query
doc = [[
    "You are a helpful, respectful and honest assistant.",
    "please keep you reply in english and do not translate into any other language in following turn!",
    "ok, I will keep my reply in english.",
    "tell me what time it is now.",
    "I'm an AI language model and don't have the ability to tell time or access real-time information.", 
    "How many characters 'c' in word 'successfully'?",
    "two 'c' ",
]]

question = ["How many characters 'c' in word 'successfully'?"]

# Generation configuration adhering to standard MT-Bench/MT-Eval protocols
generation_config = GenerationConfig(
    max_new_tokens=512,
    do_sample=False,  # Deterministic decoding for consistency
)

# Inference using REA's Heuristic Context Retrieval
res, prompts = model.generate_from_text_w_Retrieval(
    documents=doc, 
    questions=question, 
    generation_config=generation_config
)

print(f"Generated Response: {res}")