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>
121 lines
3.4 KiB
Python
121 lines
3.4 KiB
Python
"""GPU-accelerated semantic similarity for reasoning and consensus.
|
|
|
|
Provides high-level similarity operations built on the TensorBackend:
|
|
- Pairwise text similarity
|
|
- Claim deduplication with GPU cosine similarity
|
|
- Nearest-neighbor lookup for memory retrieval
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from fusionagi._logger import logger
|
|
from fusionagi.gpu.backend import TensorBackend, get_backend
|
|
|
|
|
|
def pairwise_text_similarity(
|
|
texts_a: list[str],
|
|
texts_b: list[str],
|
|
backend: TensorBackend | None = None,
|
|
) -> Any:
|
|
"""Compute pairwise cosine similarity between two sets of texts.
|
|
|
|
Args:
|
|
texts_a: First set of texts (M items).
|
|
texts_b: Second set of texts (N items).
|
|
backend: TensorBackend to use. If None, auto-selects.
|
|
|
|
Returns:
|
|
Similarity matrix of shape (M, N) as a NumPy array.
|
|
"""
|
|
be = backend or get_backend()
|
|
emb_a = be.embed_texts(texts_a)
|
|
emb_b = be.embed_texts(texts_b)
|
|
sim = be.cosine_similarity_matrix(emb_a, emb_b)
|
|
return be.to_numpy(sim)
|
|
|
|
|
|
def deduplicate_claims(
|
|
claims: list[str],
|
|
threshold: float = 0.85,
|
|
backend: TensorBackend | None = None,
|
|
) -> list[list[int]]:
|
|
"""Group semantically similar claims using GPU-accelerated similarity.
|
|
|
|
Args:
|
|
claims: List of claim texts.
|
|
threshold: Similarity threshold for grouping.
|
|
backend: TensorBackend to use.
|
|
|
|
Returns:
|
|
List of groups, where each group is a list of claim indices.
|
|
"""
|
|
if not claims:
|
|
return []
|
|
if len(claims) == 1:
|
|
return [[0]]
|
|
|
|
be = backend or get_backend()
|
|
embeddings = be.embed_texts(claims)
|
|
sim_matrix = be.to_numpy(be.cosine_similarity_matrix(embeddings, embeddings))
|
|
|
|
used: set[int] = set()
|
|
groups: list[list[int]] = []
|
|
|
|
for i in range(len(claims)):
|
|
if i in used:
|
|
continue
|
|
group = [i]
|
|
used.add(i)
|
|
for j in range(i + 1, len(claims)):
|
|
if j in used:
|
|
continue
|
|
if sim_matrix[i, j] >= threshold:
|
|
group.append(j)
|
|
used.add(j)
|
|
groups.append(group)
|
|
|
|
logger.debug(
|
|
"Claim deduplication complete",
|
|
extra={"total_claims": len(claims), "groups": len(groups)},
|
|
)
|
|
return groups
|
|
|
|
|
|
def nearest_neighbors(
|
|
query_texts: list[str],
|
|
corpus_texts: list[str],
|
|
top_k: int = 5,
|
|
backend: TensorBackend | None = None,
|
|
) -> list[list[tuple[int, float]]]:
|
|
"""Find top-k nearest neighbors from corpus for each query.
|
|
|
|
Args:
|
|
query_texts: Query texts to search for.
|
|
corpus_texts: Corpus texts to search within.
|
|
top_k: Number of nearest neighbors per query.
|
|
backend: TensorBackend to use.
|
|
|
|
Returns:
|
|
For each query, a list of (corpus_index, similarity_score) tuples.
|
|
"""
|
|
if not query_texts or not corpus_texts:
|
|
return [[] for _ in query_texts]
|
|
|
|
be = backend or get_backend()
|
|
import numpy as np
|
|
|
|
q_emb = be.embed_texts(query_texts)
|
|
c_emb = be.embed_texts(corpus_texts)
|
|
sim = be.to_numpy(be.cosine_similarity_matrix(q_emb, c_emb))
|
|
|
|
results: list[list[tuple[int, float]]] = []
|
|
for i in range(len(query_texts)):
|
|
row = sim[i]
|
|
k = min(top_k, len(corpus_texts))
|
|
top_indices = np.argsort(row)[-k:][::-1]
|
|
results.append([(int(idx), float(row[idx])) for idx in top_indices])
|
|
|
|
return results
|