import json import math import numpy as np from dataclasses import dataclass, field import os try: from sentence_transformers import SentenceTransformer _model = None except ImportError: SentenceTransformer = None _model = None def _get_model(): global _model if _model is None and SentenceTransformer is not None: _model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2') return _model def vectorize(text: str) -> list[float]: model = _get_model() if model: import logging logging.getLogger('embedder').info(f"Running inference on sentence-transformers/all-MiniLM-L6-v2 for text length {len(text)}") return model.encode([text])[0].tolist() return [] def cosine_similarity(left: list[float], right: list[float]) -> float: if not left or not right: return 0.0 dot = sum(l * r for l, r in zip(left, right)) left_norm = math.sqrt(sum(v * v for v in left)) right_norm = math.sqrt(sum(v * v for v in right)) if not left_norm or not right_norm: return 0.0 return dot / (left_norm * right_norm) @dataclass class SimpleEmbeddingIndex: entries: dict[str, list[float]] = field(default_factory=dict) def add(self, record_id: str, text: str) -> None: self.entries[record_id] = vectorize(text) def search(self, query: str, limit: int = 5) -> list[tuple[str, float]]: qvec = vectorize(query) scored = [(record_id, cosine_similarity(qvec, vec)) for record_id, vec in self.entries.items()] return sorted(scored, key=lambda item: item[1], reverse=True)[:limit] def extract_keywords(text: str, limit: int = 6) -> list[str]: # Keeping extract_keywords simple as it's not a model response import re from collections import Counter tokens = [tok.lower() for tok in re.findall(r"[A-Za-z0-9']+", text or '') if len(tok) > 2] return [word for word, _ in Counter(tokens).most_common(limit)]