From c90b48e3f8da9dd0f8d2ae82ddf977436bb0cfc3 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Sun, 15 Feb 2026 18:19:50 +0000 Subject: Initial implementation of HAG (Hopfield-Augmented Generation) Core Hopfield retrieval module with energy-based convergence guarantees, memory bank, FAISS baseline retriever, evaluation metrics, and end-to-end pipeline. All 45 tests passing on CPU with synthetic data. Co-Authored-By: Claude Opus 4.6 --- hag/datatypes.py | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 hag/datatypes.py (limited to 'hag/datatypes.py') diff --git a/hag/datatypes.py b/hag/datatypes.py new file mode 100644 index 0000000..0f4254d --- /dev/null +++ b/hag/datatypes.py @@ -0,0 +1,37 @@ +"""Data types used across HAG modules.""" + +from dataclasses import dataclass, field +from typing import List, Optional + +import torch + + +@dataclass +class HopfieldResult: + """Result from Hopfield iterative retrieval.""" + + attention_weights: torch.Tensor # (batch, N) or (N,) + converged_query: torch.Tensor # (batch, d) or (d,) + num_steps: int + trajectory: Optional[List[torch.Tensor]] = None # list of q_t + energy_curve: Optional[List[torch.Tensor]] = None # list of E(q_t) + + +@dataclass +class RetrievalResult: + """Result from a retriever (FAISS or Hopfield).""" + + passages: List[str] + scores: torch.Tensor # top-k scores + indices: torch.Tensor # top-k indices + hopfield_result: Optional[HopfieldResult] = None + + +@dataclass +class PipelineResult: + """Result from the full RAG/HAG pipeline.""" + + question: str + answer: str + retrieved_passages: List[str] + retrieval_result: RetrievalResult -- cgit v1.2.3