Memory Augmentation › A Worked Example: Document QA with Memory

class DocumentMemoryQA:
    """
    Memory-augmented document QA system.

    Processes documents in chunks, storing important information
    to memory, then retrieves relevant context for questions.
    """

    def __init__(
        self, chunk_size=100, memory_size=1000, key_dim=64, value_dim=128
    ):
        self.chunk_size = chunk_size
        self.key_dim = key_dim
        self.value_dim = value_dim

        # Memory storage
        self.memory_keys = []
        self.memory_values = []
        self.memory_texts = []  # Store original text for interpretability
        self.memory_size = memory_size

        # Simple encoder (in practice, would be a trained model)
        self.key_encoder = nn.Linear(100, key_dim)  # Simplified
        self.value_encoder = nn.Linear(100, value_dim)

    def encode_text(self, text):
        """Simple bag-of-chars encoding for demonstration."""
        # Create fixed-size representation
        chars = [ord(c) % 100 for c in text[:100]]
        chars = chars + [0] * (100 - len(chars))
        return torch.tensor(chars, dtype=torch.float32)

    def process_document(self, document):
        """Process document into memory chunks."""
        words = document.split()

        for i in range(0, len(words), self.chunk_size):
            chunk = " ".join(words[i : i + self.chunk_size])
            if len(chunk.strip()) == 0:
                continue

            # Encode chunk
            encoded = self.encode_text(chunk)
            key = self.key_encoder(encoded)
            value = self.value_encoder(encoded)

            # Store
            self.memory_keys.append(key)
            self.memory_values.append(value)
            self.memory_texts.append(chunk)

        # Truncate if needed
        if len(self.memory_keys) > self.memory_size:
            self.memory_keys = self.memory_keys[-self.memory_size :]
            self.memory_values = self.memory_values[-self.memory_size :]
            self.memory_texts = self.memory_texts[-self.memory_size :]

    def query(self, question, top_k=3):
        """Query memory for relevant context."""
        if not self.memory_keys:
            return []

        # Encode question
        q_encoded = self.encode_text(question)
        q_key = self.key_encoder(q_encoded)

        # Compute similarities
        keys_tensor = torch.stack(self.memory_keys)
        scores = torch.matmul(q_key, keys_tensor.T)

        # Get top-k
        k = min(top_k, len(self.memory_keys))
        topk_scores, topk_indices = torch.topk(scores, k)

        # Return relevant chunks
        results = []
        for idx, score in zip(topk_indices.tolist(), topk_scores.tolist()):
            results.append(
                {"chunk": self.memory_texts[idx], "score": score, "index": idx}
            )

        return results


# Create a sample document
sample_document = """
The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. 
It is named after the engineer Gustave Eiffel, whose company designed and built the tower. 
Locally nicknamed "La dame de fer" (French for "Iron Lady"), it was constructed from 1887 to 1889 
as the entrance arch for the 1889 World's Fair. The tower is 330 metres tall, about the same 
height as an 81-storey building, and the tallest structure in Paris. Its base is square, 
measuring 125 metres on each side. During its construction, the Eiffel Tower surpassed the 
Washington Monument to become the tallest human-made structure in the world, a title it held 
for 41 years until the Chrysler Building in New York City was finished in 1930.

The tower has three levels for visitors, with restaurants on the first and second levels. 
The top level's upper platform is 276 metres above the ground, the highest observation deck 
accessible to the public in the European Union. Tickets can be purchased to ascend by stairs 
or lift to the first and second levels. The climb from ground level to the first level is 
over 300 steps, as is the climb from the first level to the second. Although there is a 
staircase to the top level, it is usually accessible only by lift.
"""

# Process document and query
qa_system = DocumentMemoryQA(chunk_size=50)
qa_system.process_document(sample_document)

# Ask questions
questions = [
    "How tall is the Eiffel Tower?",
    "Who designed the tower?",
    "How many steps to climb?",
]

query_results = []
for q in questions:
    results = qa_system.query(q, top_k=2)
    query_results.append({"question": q, "results": results})