Spaces:
Running
Running
| """ | |
| Embedding-based agents using pre-computed or live embeddings. | |
| PrecomputedEmbeddingAgent: Uses pre-computed title embeddings from wiki_data | |
| LiveEmbeddingAgent: Computes embeddings on-the-fly with sentence-transformers | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from typing import TYPE_CHECKING | |
| import numpy as np | |
| from src.agents.base import Agent, AgentContext | |
| if TYPE_CHECKING: | |
| from sentence_transformers import SentenceTransformer | |
| logger = logging.getLogger(__name__) | |
| class PrecomputedEmbeddingAgent(Agent): | |
| """ | |
| Greedy agent that uses pre-computed title embeddings. | |
| Picks the link with highest cosine similarity to the target. | |
| Uses the wiki_data FAISS index for fast similarity lookup. | |
| """ | |
| def __init__(self, avoid_revisits: bool = True) -> None: | |
| """ | |
| Initialize with pre-computed embeddings. | |
| Args: | |
| avoid_revisits: If True, penalize revisiting pages | |
| """ | |
| self._avoid_revisits = avoid_revisits | |
| self._wiki_data = None # Lazy load | |
| def _ensure_loaded(self) -> None: | |
| """Lazy load wiki_data to avoid slow import.""" | |
| if self._wiki_data is None: | |
| from src.data.loader import wiki_data | |
| self._wiki_data = wiki_data | |
| def name(self) -> str: | |
| return "precomputed" | |
| def description(self) -> str: | |
| return "Greedy embedding similarity (pre-computed title embeddings)" | |
| def choose_link(self, context: AgentContext) -> str: | |
| """Pick the link most similar to target.""" | |
| self._ensure_loaded() | |
| # Always click target if available | |
| if context.target_title in context.available_links: | |
| return context.target_title | |
| # Rank by similarity to target | |
| ranked = self._wiki_data.rank_by_similarity( | |
| candidates=context.available_links, | |
| target=context.target_title, | |
| ) | |
| if not ranked: | |
| # Fallback: no embeddings found, pick first link | |
| logger.warning("No embeddings found for candidates, using fallback") | |
| return context.available_links[0] | |
| # Filter out revisits if enabled | |
| if self._avoid_revisits: | |
| visited = set(context.path_so_far) | |
| for title, _sim in ranked: | |
| if title not in visited: | |
| return title | |
| # Return best match (or first if all visited) | |
| return ranked[0][0] | |
| class LiveEmbeddingAgent(Agent): | |
| """ | |
| Greedy agent that computes embeddings on-the-fly. | |
| Uses sentence-transformers to embed article titles in real-time. | |
| More accurate than pre-computed (can use better models) but slower. | |
| """ | |
| def __init__( | |
| self, | |
| model_name: str = "all-MiniLM-L6-v2", | |
| avoid_revisits: bool = True, | |
| ) -> None: | |
| """ | |
| Initialize with a sentence-transformer model. | |
| Args: | |
| model_name: HuggingFace model name for sentence-transformers | |
| avoid_revisits: If True, penalize revisiting pages | |
| """ | |
| self._model_name = model_name | |
| self._avoid_revisits = avoid_revisits | |
| self._model: SentenceTransformer | None = None | |
| def _ensure_loaded(self) -> None: | |
| """Lazy load the model.""" | |
| if self._model is None: | |
| from sentence_transformers import SentenceTransformer | |
| logger.info(f"Loading sentence-transformer model: {self._model_name}") | |
| self._model = SentenceTransformer(self._model_name) | |
| def name(self) -> str: | |
| # Short name for the model | |
| short_name = self._model_name.split("/")[-1] | |
| return f"live-{short_name}" | |
| def description(self) -> str: | |
| return f"Greedy embedding similarity (live: {self._model_name})" | |
| def _compute_similarities( | |
| self, candidates: list[str], target: str | |
| ) -> list[tuple[str, float]]: | |
| """Compute cosine similarities between candidates and target.""" | |
| self._ensure_loaded() | |
| if not candidates: | |
| return [] | |
| # Encode target and all candidates | |
| all_texts = [target] + candidates | |
| embeddings = self._model.encode(all_texts, convert_to_numpy=True) | |
| # Normalize for cosine similarity | |
| embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True) | |
| # Target is first embedding | |
| target_emb = embeddings[0] | |
| candidate_embs = embeddings[1:] | |
| # Compute similarities | |
| similarities = np.dot(candidate_embs, target_emb) | |
| # Return sorted by similarity | |
| results = list(zip(candidates, similarities.tolist(), strict=True)) | |
| results.sort(key=lambda x: x[1], reverse=True) | |
| return [(title, sim) for title, sim in results] | |
| def choose_link(self, context: AgentContext) -> str: | |
| """Pick the link most similar to target using live embeddings.""" | |
| # Always click target if available | |
| if context.target_title in context.available_links: | |
| return context.target_title | |
| # Compute similarities | |
| ranked = self._compute_similarities( | |
| candidates=context.available_links, | |
| target=context.target_title, | |
| ) | |
| if not ranked: | |
| return context.available_links[0] | |
| # Filter out revisits if enabled | |
| if self._avoid_revisits: | |
| visited = set(context.path_so_far) | |
| for title, _sim in ranked: | |
| if title not in visited: | |
| return title | |
| return ranked[0][0] | |
| class HybridEmbeddingAgent(Agent): | |
| """ | |
| Agent that combines pre-computed and live embeddings. | |
| Uses pre-computed for quick filtering, then live for final ranking. | |
| Good balance of speed and accuracy. | |
| """ | |
| def __init__( | |
| self, | |
| model_name: str = "all-MiniLM-L6-v2", | |
| top_k: int = 20, | |
| avoid_revisits: bool = True, | |
| ) -> None: | |
| """ | |
| Initialize hybrid agent. | |
| Args: | |
| model_name: Model for live embedding | |
| top_k: Number of candidates to re-rank with live embeddings | |
| avoid_revisits: Penalize revisits | |
| """ | |
| self._precomputed = PrecomputedEmbeddingAgent(avoid_revisits=False) | |
| self._live = LiveEmbeddingAgent(model_name, avoid_revisits=False) | |
| self._top_k = top_k | |
| self._avoid_revisits = avoid_revisits | |
| def name(self) -> str: | |
| return f"hybrid-{self._top_k}" | |
| def description(self) -> str: | |
| return f"Hybrid: pre-computed filter ({self._top_k}) + live re-rank" | |
| def choose_link(self, context: AgentContext) -> str: | |
| """Two-stage selection: pre-computed filter, then live re-rank.""" | |
| # Always click target if available | |
| if context.target_title in context.available_links: | |
| return context.target_title | |
| self._precomputed._ensure_loaded() | |
| # Stage 1: Get top-k candidates from pre-computed | |
| ranked = self._precomputed._wiki_data.rank_by_similarity( | |
| candidates=context.available_links, | |
| target=context.target_title, | |
| ) | |
| if not ranked: | |
| return context.available_links[0] | |
| # Take top-k for re-ranking | |
| top_candidates = [title for title, _ in ranked[: self._top_k]] | |
| # Stage 2: Re-rank with live embeddings | |
| reranked = self._live._compute_similarities( | |
| candidates=top_candidates, | |
| target=context.target_title, | |
| ) | |
| # Filter revisits | |
| if self._avoid_revisits: | |
| visited = set(context.path_so_far) | |
| for title, _sim in reranked: | |
| if title not in visited: | |
| return title | |
| return reranked[0][0] if reranked else top_candidates[0] | |