Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import os | |
| from difflib import SequenceMatcher | |
| from typing import Any | |
| # ZeroGPU must be imported before anything that imports torch. | |
| try: | |
| import spaces # type: ignore | |
| _SPACES_AVAILABLE = True | |
| except ImportError: | |
| _SPACES_AVAILABLE = False | |
| class _Spaces: | |
| def GPU(function=None, *args, **kwargs): | |
| if function is not None and callable(function): | |
| return function | |
| def decorator(fn): | |
| return fn | |
| return decorator | |
| spaces = _Spaces() # type: ignore | |
| try: | |
| import torch | |
| except Exception: | |
| torch = None | |
| import numpy as np | |
| from memory_engine import ( | |
| build_scene, | |
| memory_context_score, | |
| normalize_text, | |
| _context_allows_memory, | |
| ) | |
| try: | |
| from sentence_transformers import SentenceTransformer | |
| except ImportError: | |
| SentenceTransformer = None | |
| try: | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| except ImportError: | |
| TfidfVectorizer = None | |
| cosine_similarity = None | |
| EMBEDDING_MODEL_NAME = os.getenv( | |
| "LEXI_EMBEDDING_MODEL", | |
| "sentence-transformers/all-MiniLM-L6-v2", | |
| ) | |
| USE_ZEROGPU = os.getenv("LEXI_USE_ZEROGPU", "1") != "0" | |
| _MINILM_ERROR = "" | |
| if torch is not None and (_SPACES_AVAILABLE or torch.cuda.is_available()): | |
| DEVICE = "cuda" if USE_ZEROGPU else "cpu" | |
| else: | |
| DEVICE = "cpu" | |
| _EMBEDDER = None | |
| if SentenceTransformer is not None: | |
| try: | |
| _EMBEDDER = SentenceTransformer( | |
| EMBEDDING_MODEL_NAME, | |
| device=DEVICE, | |
| ) | |
| except Exception as exc: | |
| _MINILM_ERROR = f"MiniLM initialization failed: {exc}" | |
| _EMBEDDER = None | |
| else: | |
| _MINILM_ERROR = "sentence-transformers is not installed." | |
| def _encode_gpu(texts: list[str]) -> np.ndarray: | |
| if _EMBEDDER is None: | |
| raise RuntimeError(_MINILM_ERROR or "MiniLM model is unavailable.") | |
| return np.asarray( | |
| _EMBEDDER.encode( | |
| texts, | |
| normalize_embeddings=True, | |
| convert_to_numpy=True, | |
| show_progress_bar=False, | |
| batch_size=min(32, max(1, len(texts))), | |
| ), | |
| dtype=np.float32, | |
| ) | |
| _PROFILE_CACHE: dict[tuple[str, str], np.ndarray] = {} | |
| _CACHE_LIMIT = 512 | |
| def _encode(texts: list[str]) -> np.ndarray: | |
| if not texts: | |
| return np.empty((0, 384), dtype=np.float32) | |
| if _EMBEDDER is None: | |
| raise RuntimeError(_MINILM_ERROR or "MiniLM model is unavailable.") | |
| return _encode_gpu(texts) | |
| def _candidate_retrieval_pool( | |
| memories: list[dict[str, Any]], | |
| ) -> list[dict[str, Any]]: | |
| pool = [] | |
| for memory in memories or []: | |
| status = str(memory.get("status", "candidate") or "candidate") | |
| if status in {"active", "candidate", "ambiguous"}: | |
| confidence = float(memory.get("confidence", 0.0) or 0.0) | |
| positives = int(memory.get("positive_evidence_count", 0) or 0) | |
| if status == "active" or (confidence >= 0.60 and positives >= 1): | |
| pool.append(memory) | |
| return pool | |
| def retrieval_memories(memories: list[dict[str, Any]]) -> list[dict[str, Any]]: | |
| return _candidate_retrieval_pool(memories) | |
| def _memory_profile(memory: dict[str, Any]) -> str: | |
| positives = memory.get("positive_contexts", [])[-8:] | |
| negatives = memory.get("negative_contexts", [])[-8:] | |
| contexts = " ".join(item.get("context", "") for item in positives) | |
| negative_contexts = " ".join(item.get("context", "") for item in negatives) | |
| tags = " ".join( | |
| tag | |
| for item in positives | |
| for tag in item.get("tags", []) | |
| ) | |
| return " ".join( | |
| part | |
| for part in [ | |
| memory.get("source", ""), | |
| memory.get("target", ""), | |
| memory.get("kind", ""), | |
| "preferred contexts", | |
| contexts, | |
| "avoid contexts", | |
| negative_contexts, | |
| tags, | |
| ] | |
| if part | |
| )[:3500] | |
| def _lexical_score(query: str, memory: dict[str, Any]) -> float: | |
| query_norm = normalize_text(query) | |
| best = 0.0 | |
| for term in [memory.get("source", ""), *memory.get("aliases", [])]: | |
| term_norm = normalize_text(term) | |
| if not term_norm: | |
| continue | |
| if term_norm in query_norm: | |
| return 1.0 | |
| for token in query_norm.split(): | |
| if len(token) >= 3 and len(term_norm) >= 3: | |
| best = max( | |
| best, | |
| SequenceMatcher( | |
| None, | |
| token, | |
| term_norm, | |
| autojunk=False, | |
| ).ratio(), | |
| ) | |
| return min(best, 0.94) | |
| def _term_present(query: str, memory: dict[str, Any]) -> bool: | |
| query_norm = normalize_text(query) | |
| target_norm = normalize_text(memory.get("target", "")) | |
| for term in [memory.get("source", ""), *memory.get("aliases", [])]: | |
| term_norm = normalize_text(term) | |
| if not term_norm: | |
| continue | |
| if term_norm == target_norm and memory.get("kind") != "casing": | |
| continue | |
| if term_norm in query_norm: | |
| return True | |
| return False | |
| def _context_features( | |
| query: str, | |
| memory: dict[str, Any], | |
| note_title: str, | |
| scene_text: str, | |
| ) -> dict[str, Any]: | |
| scene = build_scene(note_title, scene_text or "", query) | |
| ctx = memory_context_score(memory, scene) | |
| positive = float(ctx["positive"]) | |
| negative = float(ctx["negative"]) | |
| context_fit = max(0.0, positive - 0.65 * negative) | |
| blocked = ( | |
| not _context_allows_memory(memory, scene, positive, negative) | |
| or ctx["exception"] >= 0.76 | |
| or (negative >= 0.58 and negative > positive) | |
| ) | |
| return { | |
| "scene": scene, | |
| "positive": positive, | |
| "negative": negative, | |
| "exception": float(ctx["exception"]), | |
| "context_fit": context_fit, | |
| "blocked": blocked, | |
| } | |
| def _quick_results( | |
| query: str, | |
| memories: list[dict], | |
| note_title: str, | |
| scene_text: str, | |
| top_k: int, | |
| ) -> list[dict[str, Any]]: | |
| results = [] | |
| for memory in memories: | |
| features = _context_features(query, memory, note_title, scene_text) | |
| lexical = _lexical_score(query, memory) | |
| context_fit = features["context_fit"] | |
| evidence = min( | |
| 1.0, | |
| float(memory.get("positive_evidence_count", 0) or 0) / 3.0, | |
| ) | |
| score = ( | |
| 0.47 * lexical | |
| + 0.28 * context_fit | |
| + 0.15 * float(memory.get("confidence", 0.0) or 0.0) | |
| + 0.10 * evidence | |
| ) | |
| if features["blocked"]: | |
| score *= 0.20 | |
| results.append({ | |
| "memory_id": memory["id"], | |
| "source": memory["source"], | |
| "target": memory["target"], | |
| "kind": memory["kind"], | |
| "status": memory.get("status", "candidate"), | |
| "confidence": round(float(memory.get("confidence", 0.0) or 0.0), 3), | |
| "evidence_count": memory.get("evidence_count", 0), | |
| "semantic_score": None, | |
| "lexical_score": round(lexical, 4), | |
| "context_score": round(context_fit, 4), | |
| "positive_context_score": round(features["positive"], 4), | |
| "negative_context_score": round(features["negative"], 4), | |
| "retrieval_score": round(score, 4), | |
| "retrieval_backend": "fast lexical + context", | |
| "blocked_by_context": bool(features["blocked"]), | |
| "term_present": _term_present(query, memory), | |
| "reason": "Fast lexical/context retrieval.", | |
| }) | |
| results.sort(key=lambda x: x["retrieval_score"], reverse=True) | |
| return results[:top_k] | |
| def quick_retrieve_memories( | |
| query: str, | |
| memories: list[dict], | |
| top_k: int = 4, | |
| fuzzy_threshold: float = 0.96, | |
| note_title: str = "", | |
| scene_text: str = "", | |
| ): | |
| pool = retrieval_memories(memories) | |
| if not pool or not query.strip(): | |
| return [], True | |
| results = _quick_results( | |
| query, | |
| pool, | |
| note_title, | |
| scene_text, | |
| max(top_k, 8), | |
| ) | |
| strong = [ | |
| item | |
| for item in results | |
| if item["term_present"] | |
| and item.get("status") == "active" | |
| and int(item.get("evidence_count", 0) or 0) >= 2 | |
| and not item["blocked_by_context"] | |
| and item["lexical_score"] >= fuzzy_threshold | |
| and ( | |
| item["context_score"] >= 0.45 | |
| or item["negative_context_score"] < 0.50 | |
| ) | |
| ] | |
| unresolved_candidate = any( | |
| item.get("status") in {"candidate", "ambiguous"} | |
| and not item.get("blocked_by_context") | |
| and ( | |
| item.get("term_present") | |
| or item.get("lexical_score", 0.0) >= 0.45 | |
| or item.get("context_score", 0.0) >= 0.45 | |
| ) | |
| for item in results | |
| ) | |
| if strong and not unresolved_candidate: | |
| return strong[:top_k], True | |
| return results[:top_k], False | |
| def _statistical_batch_scores( | |
| query: str, | |
| memories: list[dict[str, Any]], | |
| ) -> tuple[np.ndarray, str]: | |
| if TfidfVectorizer is None or cosine_similarity is None: | |
| return np.zeros(len(memories), dtype=np.float32), "ML unavailable" | |
| texts = [_memory_profile(memory) for memory in memories] | |
| corpus = [query, *texts] | |
| word_vectorizer = TfidfVectorizer( | |
| lowercase=True, | |
| ngram_range=(1, 2), | |
| min_df=1, | |
| sublinear_tf=True, | |
| ) | |
| char_vectorizer = TfidfVectorizer( | |
| lowercase=True, | |
| analyzer="char_wb", | |
| ngram_range=(2, 5), | |
| min_df=1, | |
| sublinear_tf=True, | |
| ) | |
| word_matrix = word_vectorizer.fit_transform(corpus) | |
| char_matrix = char_vectorizer.fit_transform(corpus) | |
| word_scores = cosine_similarity(word_matrix[0:1], word_matrix[1:]).ravel() | |
| char_scores = cosine_similarity(char_matrix[0:1], char_matrix[1:]).ravel() | |
| scores = 0.40 * word_scores + 0.60 * char_scores | |
| return np.asarray(scores, dtype=np.float32), "ML fallback: word+char TF-IDF" | |
| def _semantic_batch_scores( | |
| query: str, | |
| memories: list[dict[str, Any]], | |
| ) -> tuple[np.ndarray, str]: | |
| profiles = [_memory_profile(memory) for memory in memories] | |
| missing = [] | |
| missing_keys = [] | |
| vectors = [None] * len(memories) | |
| for index, (memory, profile) in enumerate(zip(memories, profiles)): | |
| key = ( | |
| memory["id"], | |
| f'{memory.get("updated_at", "")}::{profile}', | |
| ) | |
| cached = _PROFILE_CACHE.get(key) | |
| if cached is None: | |
| missing.append(profile) | |
| missing_keys.append((index, key)) | |
| else: | |
| vectors[index] = cached | |
| encoded = _encode([query, *missing]) | |
| query_vector = encoded[0] | |
| for offset, (index, key) in enumerate(missing_keys, start=1): | |
| vectors[index] = encoded[offset] | |
| _PROFILE_CACHE[key] = encoded[offset] | |
| while len(_PROFILE_CACHE) > _CACHE_LIMIT: | |
| _PROFILE_CACHE.pop(next(iter(_PROFILE_CACHE))) | |
| scores = np.asarray( | |
| [float(np.dot(vector, query_vector)) for vector in vectors], | |
| dtype=np.float32, | |
| ) | |
| return scores, "MiniLM via ZeroGPU" if DEVICE == "cuda" else "MiniLM on CPU" | |
| def retrieve_memories( | |
| query: str, | |
| memories: list[dict], | |
| note_title: str = "", | |
| scene_text: str = "", | |
| top_k: int = 4, | |
| ) -> list[dict]: | |
| pool = retrieval_memories(memories) | |
| if not pool or not query.strip(): | |
| return [] | |
| quick, strong = quick_retrieve_memories( | |
| query, | |
| pool, | |
| top_k=max(top_k, 8), | |
| note_title=note_title, | |
| scene_text=scene_text, | |
| ) | |
| if strong: | |
| return quick[:top_k] | |
| scene = build_scene(note_title, scene_text or "", query) | |
| ml_error = "" | |
| try: | |
| semantic_scores, backend = _semantic_batch_scores( | |
| scene["text"][:3000], | |
| pool, | |
| ) | |
| except Exception as exc: | |
| ml_error = str(exc) | |
| semantic_scores, backend = _statistical_batch_scores( | |
| scene["text"][:3000], | |
| pool, | |
| ) | |
| if backend == "ML unavailable": | |
| return quick[:top_k] | |
| results = [] | |
| for index, memory in enumerate(pool): | |
| features = _context_features(query, memory, note_title, scene_text) | |
| semantic = max(0.0, min(1.0, float(semantic_scores[index]))) | |
| lexical = _lexical_score(query, memory) | |
| context_fit = features["context_fit"] | |
| evidence = min( | |
| 1.0, | |
| float(memory.get("positive_evidence_count", 0) or 0) / 3.0, | |
| ) | |
| score = ( | |
| 0.42 * semantic | |
| + 0.23 * lexical | |
| + 0.25 * context_fit | |
| + 0.10 * float(memory.get("confidence", 0.0) or 0.0) | |
| ) | |
| if features["blocked"]: | |
| score *= 0.20 | |
| results.append({ | |
| "memory_id": memory["id"], | |
| "source": memory["source"], | |
| "target": memory["target"], | |
| "kind": memory["kind"], | |
| "status": memory.get("status", "candidate"), | |
| "confidence": round(float(memory.get("confidence", 0.0) or 0.0), 3), | |
| "evidence_count": memory.get("evidence_count", 0), | |
| "semantic_score": round(semantic, 4), | |
| "lexical_score": round(lexical, 4), | |
| "context_score": round(context_fit, 4), | |
| "positive_context_score": round(features["positive"], 4), | |
| "negative_context_score": round(features["negative"], 4), | |
| "retrieval_score": round(score, 4), | |
| "retrieval_backend": backend, | |
| "blocked_by_context": bool(features["blocked"]), | |
| "term_present": _term_present(query, memory), | |
| "reason": ( | |
| "MiniLM semantic retrieval validated the scene." | |
| if backend.startswith("MiniLM") | |
| else "Statistical ML fallback handled semantic/phonetic retrieval." | |
| ), | |
| "ml_error": ml_error if ml_error else None, | |
| }) | |
| results.sort(key=lambda x: x["retrieval_score"], reverse=True) | |
| return results[:top_k] | |
| def ml_probe() -> dict[str, Any]: | |
| started = __import__("time").perf_counter() | |
| result = { | |
| "ok": False, | |
| "backend": backend_status(), | |
| "model": EMBEDDING_MODEL_NAME, | |
| "device": DEVICE, | |
| "minilm_loaded": _EMBEDDER is not None, | |
| "error": None, | |
| } | |
| try: | |
| if _EMBEDDER is None: | |
| result["backend"] = "ML fallback: word+char TF-IDF" | |
| vector_scores, backend = _statistical_batch_scores( | |
| "Lexi ML health check", | |
| [normalize_memory_for_probe()], | |
| ) | |
| result["backend"] = backend | |
| result["ok"] = len(vector_scores) == 1 | |
| else: | |
| vectors = _encode(["Lexi ML health check"]) | |
| result["backend"] = "MiniLM via ZeroGPU" if DEVICE == "cuda" else "MiniLM on CPU" | |
| result["ok"] = bool(vectors.shape[0] == 1) | |
| except Exception as exc: | |
| result["error"] = str(exc) | |
| result["latency_ms"] = round((__import__("time").perf_counter() - started) * 1000, 1) | |
| return result | |
| def normalize_memory_for_probe() -> dict[str, Any]: | |
| return { | |
| "id": "ml_probe", | |
| "source": "lexi", | |
| "target": "Lexi", | |
| "kind": "product_service", | |
| "confidence": 0.9, | |
| "positive_evidence_count": 1, | |
| "evidence_count": 1, | |
| "positive_contexts": [{"context": "Lexi personal writing assistant", "tags": ["software"]}], | |
| "negative_contexts": [], | |
| "exceptions": [], | |
| } | |
| def backend_status() -> str: | |
| if _EMBEDDER is not None: | |
| if DEVICE == "cuda": | |
| return "MiniLM semantic retrieval via ZeroGPU (loaded)" | |
| return "MiniLM semantic retrieval on CPU (loaded)" | |
| if _MINILM_ERROR: | |
| return f"ML fallback enabled; {_MINILM_ERROR}" | |
| return "ML fallback enabled" | |
| def ml_diagnostics() -> dict[str, Any]: | |
| return { | |
| "minilm_loaded": _EMBEDDER is not None, | |
| "model": EMBEDDING_MODEL_NAME, | |
| "device": DEVICE, | |
| "zero_gpu_requested": USE_ZEROGPU, | |
| "spaces_available": _SPACES_AVAILABLE, | |
| "minilm_error": _MINILM_ERROR, | |
| "fallback_available": TfidfVectorizer is not None, | |
| "backend": backend_status(), | |
| } | |