feat: GPU/TensorCore integration — TensorFlow backend, GPU-accelerated reasoning, training, and memory
Some checks failed
Some checks failed
- New fusionagi/gpu/ module with TensorBackend protocol abstraction - TensorFlowBackend: GPU-accelerated ops with TensorCore mixed-precision - NumPyBackend: CPU fallback (always available, no extra deps) - Auto-selects best available backend at runtime - GPU-accelerated operations: - Cosine similarity matrix (batched, XLA-compiled) - Multi-head attention for consensus scoring - Batch hypothesis scoring on GPU - Semantic similarity search (pairwise, nearest-neighbor, deduplication) - New TensorFlowAdapter (fusionagi/adapters/): - LLMAdapter for local TF/Keras model inference - TensorCore mixed-precision support - GPU-accelerated embedding synthesis fallback - Reasoning pipeline integration: - gpu_scoring.py: drop-in GPU replacement for multi_path scoring - Super Big Brain: use_gpu config flag, GPU scoring when available - Memory integration: - gpu_search.py: GPU-accelerated semantic search for SemanticGraphMemory - Self-improvement integration: - gpu_training.py: gradient-based heuristic weight optimization - Reflective memory training loop with loss tracking - Dependencies: gpu extra (tensorflow>=2.16, numpy>=1.26) - 64 new tests (276 total), all passing - Architecture spec: docs/gpu_tensorcore_integration.md Co-Authored-By: Nakamoto, S <defi@defi-oracle.io>
This commit is contained in:
86
fusionagi/memory/gpu_search.py
Normal file
86
fusionagi/memory/gpu_search.py
Normal file
@@ -0,0 +1,86 @@
|
||||
"""GPU-accelerated semantic search for memory subsystems.
|
||||
|
||||
Provides vector similarity search using GPU-accelerated embeddings
|
||||
for SemanticGraphMemory and EpisodicMemory.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fusionagi._logger import logger
|
||||
from fusionagi.schemas.atomic import AtomicSemanticUnit
|
||||
|
||||
|
||||
def semantic_search(
|
||||
query: str,
|
||||
units: list[AtomicSemanticUnit],
|
||||
top_k: int = 10,
|
||||
) -> list[tuple[AtomicSemanticUnit, float]]:
|
||||
"""Search atomic semantic units by vector similarity using GPU.
|
||||
|
||||
Args:
|
||||
query: Query text to search for.
|
||||
units: List of atomic semantic units to search within.
|
||||
top_k: Number of top results to return.
|
||||
|
||||
Returns:
|
||||
List of (unit, similarity_score) tuples sorted by score descending.
|
||||
"""
|
||||
if not units:
|
||||
return []
|
||||
|
||||
try:
|
||||
from fusionagi.gpu.tensor_similarity import nearest_neighbors
|
||||
|
||||
corpus = [u.content for u in units]
|
||||
results = nearest_neighbors([query], corpus, top_k=top_k)
|
||||
if not results or not results[0]:
|
||||
return []
|
||||
|
||||
return [(units[idx], score) for idx, score in results[0] if idx < len(units)]
|
||||
except ImportError:
|
||||
return _cpu_fallback_search(query, units, top_k)
|
||||
|
||||
|
||||
def _cpu_fallback_search(
|
||||
query: str,
|
||||
units: list[AtomicSemanticUnit],
|
||||
top_k: int,
|
||||
) -> list[tuple[AtomicSemanticUnit, float]]:
|
||||
"""CPU fallback: simple word-overlap similarity."""
|
||||
query_words = set(query.lower().split())
|
||||
scored: list[tuple[AtomicSemanticUnit, float]] = []
|
||||
|
||||
for unit in units:
|
||||
unit_words = set(unit.content.lower().split())
|
||||
if not unit_words:
|
||||
continue
|
||||
overlap = len(query_words & unit_words)
|
||||
score = overlap / max(len(query_words | unit_words), 1)
|
||||
scored.append((unit, score))
|
||||
|
||||
scored.sort(key=lambda x: x[1], reverse=True)
|
||||
return scored[:top_k]
|
||||
|
||||
|
||||
def batch_embed_units(
|
||||
units: list[AtomicSemanticUnit],
|
||||
) -> Any:
|
||||
"""Embed a batch of atomic semantic units using GPU.
|
||||
|
||||
Args:
|
||||
units: Units to embed.
|
||||
|
||||
Returns:
|
||||
Embedding tensor (backend-specific type).
|
||||
"""
|
||||
try:
|
||||
from fusionagi.gpu.backend import get_backend
|
||||
|
||||
be = get_backend()
|
||||
texts = [u.content for u in units]
|
||||
return be.embed_texts(texts)
|
||||
except ImportError:
|
||||
logger.debug("GPU not available for batch embedding")
|
||||
return None
|
||||
Reference in New Issue
Block a user