my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	"""
#2	Mnemosyne Dense Retrieval
#3	Supports local fastembed (ONNX) and OpenAI-compatible API embeddings.
#4	Falls back to keyword-only if neither is available.
#5	"""
#6	from __future__ import annotations
#7
#8	import json
#9	import os
#10	import urllib.request
#11	from typing import List, Optional
#12	from functools import lru_cache
#13
#14	try:
#15	import numpy as np
#16	except ImportError:
#17	np = None
#18
#19	# --- fastembed (local ONNX) ---
#20	try:
#21	from fastembed import TextEmbedding
#22	except Exception:
#23	TextEmbedding = None
#24
#25	_FASTEMBED_AVAILABLE = np is not None and TextEmbedding is not None
#26	_FASTEMBED_CACHE_DIR = os.path.join(os.path.expanduser("~/.hermes"), "cache", "fastembed")
#27
#28	# --- OpenAI-compatible API ---
#29	_OPENAI_API_KEY = os.environ.get("OPENROUTER_API_KEY", os.environ.get("OPENAI_API_KEY", ""))
#30	_OPENAI_BASE_URL = os.environ.get("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
#31
#32	# --- Model selection ---
#33	_DEFAULT_MODEL = os.environ.get("MNEMOSYNE_EMBEDDING_MODEL", "BAAI/bge-small-en-v1.5")
#34	_embedding_model = None
#35	_API_CALL_COUNT = 0
#36
#37
#38	def _is_api_model(model_name: str) -> bool:
#39	"""Check if the model should use the OpenAI-compatible API."""
#40	return (
#41	model_name.startswith("openai/")
#42	or "text-embedding" in model_name
#43	or model_name.startswith("text-embedding")
#44	)
#45
#46
#47	def _get_embedding_dim(model_name: str) -> int:
#48	"""Return the embedding dimension for a given model."""
#49	dims = {
#50	"BAAI/bge-small-en-v1.5": 384,
#51	"BAAI/bge-base-en-v1.5": 768,
#52	"BAAI/bge-large-en-v1.5": 1024,
#53	"openai/text-embedding-3-small": 1536,
#54	"openai/text-embedding-3-large": 3072,
#55	"text-embedding-3-small": 1536,
#56	"text-embedding-3-large": 3072,
#57	}
#58	return dims.get(model_name, 384)
#59
#60
#61	def _get_model():
#62	"""Lazy-load the embedding model (local fastembed)."""
#63	global _embedding_model
#64	if _is_api_model(_DEFAULT_MODEL):
#65	return "api" # Sentinel for API mode
#66	if not _FASTEMBED_AVAILABLE:
#67	return None
#68	if _embedding_model is None:
#69	os.makedirs(_FASTEMBED_CACHE_DIR, exist_ok=True)
#70	_embedding_model = TextEmbedding(
#71	model_name=_DEFAULT_MODEL,
#72	cache_dir=_FASTEMBED_CACHE_DIR,
#73	)
#74	return _embedding_model
#75
#76
#77	def _embed_api(texts: List[str]) -> Optional[np.ndarray]:
#78	"""Embed texts via OpenAI-compatible API (OpenRouter)."""
#79	global _API_CALL_COUNT
#80	if not _OPENAI_API_KEY:
#81	return None
#82
#83	url = f"{_OPENAI_BASE_URL.rstrip('/')}/embeddings"
#84	payload = json.dumps({
#85	"model": _DEFAULT_MODEL,
#86	"input": texts,
#87	}).encode()
#88
#89	headers = {
#90	"Authorization": f"Bearer {_OPENAI_API_KEY}",
#91	"Content-Type": "application/json",
#92	"HTTP-Referer": "https://mnemosyne.site",
#93	"X-Title": "Mnemosyne Embedding",
#94	}
#95
#96	for attempt in range(3):
#97	try:
#98	req = urllib.request.Request(url, data=payload, headers=headers)
#99	with urllib.request.urlopen(req, timeout=30) as resp:
#100	data = json.loads(resp.read())
#101	embeddings = [item["embedding"] for item in data["data"]]
#102	_API_CALL_COUNT += 1
#103	return np.array(embeddings, dtype=np.float32)
#104	except Exception as e:
#105	if "429" in str(e) or "rate" in str(e).lower():
#106	import time
#107	time.sleep(2 ** attempt)
#108	continue
#109	return None
#110
#111	return None
#112
#113
#114	def available() -> bool:
#115	"""Check if dense retrieval is available."""
#116	if _is_api_model(_DEFAULT_MODEL):
#117	return bool(_OPENAI_API_KEY)
#118	return _FASTEMBED_AVAILABLE and _get_model() is not None
#119
#120
#121	def available_api() -> bool:
#122	"""Check if API-based embeddings are available."""
#123	return bool(_OPENAI_API_KEY)
#124
#125
#126	@lru_cache(maxsize=512)
#127	def embed_query(text: str) -> Optional[np.ndarray]:
#128	"""Encode a single query text into a dense vector."""
#129	if not text:
#130	return None
#131
#132	if _is_api_model(_DEFAULT_MODEL):
#133	result = _embed_api([text])
#134	return result[0] if result is not None else None
#135
#136	model = _get_model()
#137	if model is None or model == "api":
#138	return None
#139	vectors = list(model.embed([text]))
#140	if not vectors:
#141	return None
#142	return vectors[0].astype(np.float32)
#143
#144
#145	def embed(texts: List[str]) -> Optional[np.ndarray]:
#146	"""Encode texts into dense vectors."""
#147	if not texts:
#148	return None
#149
#150	if _is_api_model(_DEFAULT_MODEL):
#151	return _embed_api(texts)
#152
#153	# Use cached single-query path for common case of 1 text
#154	if len(texts) == 1:
#155	v = embed_query(texts[0])
#156	if v is None:
#157	return None
#158	return np.stack([v])
#159
#160	model = _get_model()
#161	if model is None or model == "api":
#162	return None
#163	vectors = list(model.embed(texts))
#164	return np.stack(vectors).astype(np.float32)
#165
#166
#167	def serialize(vec: np.ndarray) -> str:
#168	"""Serialize embedding to JSON string."""
#169	return json.dumps(vec.tolist())
#170
#171
#172	# Export dimension for other modules
#173	EMBEDDING_DIM = _get_embedding_dim(_DEFAULT_MODEL)
#174	_DEFAULT_MODEL = _DEFAULT_MODEL # Re-export for beam.py
#175

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public