Spaces:
Runtime error
Runtime error
File size: 4,714 Bytes
9c959a8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
import sqlite3
import numpy as np
from typing import Optional, List
from pathlib import Path
import time
class EmbeddingCache:
def __init__(self, db_path: str = "embeddings_cache.db"):
"""
Ініціалізація кешу ембедінгів
Args:
db_path: шлях до файлу SQLite бази даних
"""
self.db_path = db_path
self._init_db()
self.hits = 0
self.misses = 0
def _init_db(self):
"""Ініціалізація структури бази даних"""
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS embeddings (
text_hash TEXT PRIMARY KEY,
text TEXT NOT NULL,
model TEXT NOT NULL,
embedding BLOB NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
# Індекс для швидкого пошуку за хешем
conn.execute("""
CREATE INDEX IF NOT EXISTS idx_text_hash
ON embeddings(text_hash)
""")
def _get_hash(self, text: str, model: str) -> str:
"""Створення унікального хешу для тексту та моделі"""
return str(hash(f"{text}:{model}"))
def get(self, text: str, model: str) -> Optional[np.ndarray]:
"""
Отримання ембедінгу з кешу
Args:
text: текст для пошуку
model: назва моделі ембедінгів
Returns:
np.ndarray якщо знайдено, None якщо не знайдено
"""
text_hash = self._get_hash(text, model)
with sqlite3.connect(self.db_path) as conn:
result = conn.execute(
"SELECT embedding FROM embeddings WHERE text_hash = ?",
(text_hash,)
).fetchone()
if result:
self.hits += 1
return np.frombuffer(result[0], dtype=np.float32)
self.misses += 1
return None
def put(self, text: str, model: str, embedding: np.ndarray) -> None:
"""
Збереження ембедінгу в кеш
Args:
text: вхідний текст
model: назва моделі
embedding: ембедінг для збереження
"""
text_hash = self._get_hash(text, model)
with sqlite3.connect(self.db_path) as conn:
conn.execute(
"""
INSERT OR REPLACE INTO embeddings
(text_hash, text, model, embedding)
VALUES (?, ?, ?, ?)
""",
(
text_hash,
text,
model,
np.array(embedding, dtype=np.float32).tobytes()
)
)
def clear_old(self, days: int = 30) -> int:
"""
Очищення старих записів з кешу
Args:
days: кількість днів, старіші записи будуть видалені
Returns:
Кількість видалених записів
"""
with sqlite3.connect(self.db_path) as conn:
cursor = conn.execute(
"""
DELETE FROM embeddings
WHERE created_at < datetime('now', ?)
""",
(f"-{days} days",)
)
return cursor.rowcount
def get_stats(self) -> dict:
"""Отримання статистики використання кешу"""
with sqlite3.connect(self.db_path) as conn:
total = conn.execute(
"SELECT COUNT(*) FROM embeddings"
).fetchone()[0]
size = Path(self.db_path).stat().st_size / (1024 * 1024) # Size in MB
if self.hits + self.misses > 0:
hit_rate = self.hits / (self.hits + self.misses) * 100
else:
hit_rate = 0
return {
"total_entries": total,
"cache_size_mb": round(size, 2),
"hits": self.hits,
"misses": self.misses,
"hit_rate_percent": round(hit_rate, 2)
} |