Замена RuBERT на sentence-transformers/all-MiniLM-L12-v2, упрощение формулы расчета, поддержка загрузки из отдельных .npy файлов
This commit is contained in:
@@ -27,19 +27,17 @@ class VectorStore:
|
||||
примеры. Использует косинусное сходство для расчета скора.
|
||||
|
||||
Attributes:
|
||||
vector_dim: Размерность векторов (768 для ruBERT)
|
||||
vector_dim: Размерность векторов (384 для all-MiniLM-L12-v2)
|
||||
max_examples: Максимальное количество примеров каждого типа
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector_dim: int = 768,
|
||||
vector_dim: int = 384,
|
||||
max_examples: int = 10000,
|
||||
storage_path: Optional[str] = None,
|
||||
score_multiplier: float = 5.0,
|
||||
k_min: int = 5,
|
||||
k_max: int = 10,
|
||||
base_multiplier_factor: float = 15.0,
|
||||
k: int = 3,
|
||||
):
|
||||
"""
|
||||
Инициализация хранилища.
|
||||
@@ -48,18 +46,14 @@ class VectorStore:
|
||||
vector_dim: Размерность векторов
|
||||
max_examples: Максимальное количество примеров каждого типа
|
||||
storage_path: Путь для сохранения/загрузки векторов (опционально)
|
||||
score_multiplier: Базовый множитель для усиления разницы в скорах
|
||||
k_min: Минимальное значение k для топ-k ближайших примеров
|
||||
k_max: Максимальное значение k для топ-k ближайших примеров
|
||||
base_multiplier_factor: Множитель для базового score_multiplier
|
||||
score_multiplier: Множитель для масштабирования разницы в скорах
|
||||
k: Количество ближайших примеров для расчета среднего сходства
|
||||
"""
|
||||
self.vector_dim = vector_dim
|
||||
self.max_examples = max_examples
|
||||
self.storage_path = storage_path
|
||||
self.score_multiplier = score_multiplier
|
||||
self.k_min = k_min
|
||||
self.k_max = k_max
|
||||
self.base_multiplier_factor = base_multiplier_factor
|
||||
self.k = k
|
||||
|
||||
# Инициализируем пустые массивы
|
||||
# Используем список для динамического добавления, потом конвертируем в numpy
|
||||
@@ -72,8 +66,15 @@ class VectorStore:
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# Пытаемся загрузить сохраненные векторы
|
||||
if storage_path and os.path.exists(storage_path):
|
||||
self._load_from_disk()
|
||||
# Проверяем наличие storage_path или отдельных .npy файлов
|
||||
if storage_path:
|
||||
storage_dir = Path(storage_path).parent
|
||||
positive_npy = storage_dir / "positive_embeddings.npy"
|
||||
negative_npy = storage_dir / "negative_embeddings.npy"
|
||||
|
||||
# Загружаем если есть .npz файл или отдельные .npy файлы
|
||||
if os.path.exists(storage_path) or positive_npy.exists() or negative_npy.exists():
|
||||
self._load_from_disk()
|
||||
|
||||
@property
|
||||
def positive_count(self) -> int:
|
||||
@@ -292,35 +293,23 @@ class VectorStore:
|
||||
else:
|
||||
neg_similarities = np.array([])
|
||||
|
||||
# Используем топ-k ближайших примеров для более чувствительной оценки
|
||||
# Берем k в диапазоне [k_min, k_max] для большей чувствительности к различиям
|
||||
k_pos = min(self.k_max, max(self.k_min, len(pos_similarities)))
|
||||
|
||||
# Топ-k положительных примеров (самые близкие)
|
||||
top_k_pos_sim = float(np.mean(np.sort(pos_similarities)[-k_pos:]))
|
||||
# Используем топ-k ближайших примеров для расчета среднего сходства
|
||||
k_pos = min(self.k, len(pos_similarities))
|
||||
top_k_pos = np.sort(pos_similarities)[-k_pos:]
|
||||
avg_pos = float(np.mean(top_k_pos))
|
||||
|
||||
# Для отрицательных: если их меньше k, берем все, иначе топ-k
|
||||
if len(neg_similarities) > 0:
|
||||
k_neg = min(self.k_max, max(self.k_min, len(neg_similarities)))
|
||||
top_k_neg_sim = float(np.mean(np.sort(neg_similarities)[-k_neg:]))
|
||||
k_neg = min(self.k, len(neg_similarities))
|
||||
top_k_neg = np.sort(neg_similarities)[-k_neg:]
|
||||
avg_neg = float(np.mean(top_k_neg))
|
||||
else:
|
||||
# Если нет отрицательных примеров, используем нейтральное значение
|
||||
top_k_neg_sim = top_k_pos_sim # Нейтральный скор = 0.5
|
||||
avg_neg = avg_pos # Нейтральный скор = 0.5
|
||||
|
||||
# === Вариант 1: neg/pos (разница между топ-k положительными и отрицательными) ===
|
||||
# Используем более агрессивную нормализацию для малых различий
|
||||
diff = top_k_pos_sim - top_k_neg_sim
|
||||
|
||||
# Увеличиваем множитель для большей чувствительности к малым различиям
|
||||
# Базовый множитель умножаем на base_multiplier_factor для работы с топ-k
|
||||
base_multiplier = self.score_multiplier * self.base_multiplier_factor
|
||||
|
||||
# Адаптивный множитель: чем больше примеров, тем выше чувствительность
|
||||
# При 500 примерах: 1.25, при 1000+: 1.5
|
||||
adaptive_multiplier = base_multiplier * (1.0 + min(0.5, (self.positive_count + self.negative_count) / 2000))
|
||||
|
||||
score_neg_pos = 0.5 + (diff * adaptive_multiplier)
|
||||
score_neg_pos = max(0.0, min(1.0, score_neg_pos))
|
||||
# Формула расчета score: (diff * scale + 1) / 2, переводим из [-1, 1] в [0, 1]
|
||||
diff = avg_pos - avg_neg
|
||||
score_neg_pos = np.clip((diff * self.score_multiplier + 1) / 2, 0.0, 1.0)
|
||||
|
||||
# === Вариант 2: pos only (только положительные, топ-k ближайших) ===
|
||||
# Берём топ-5 ближайших положительных примеров
|
||||
@@ -352,9 +341,9 @@ class VectorStore:
|
||||
neg_mean = neg_std = neg_min = neg_max = 0.0
|
||||
|
||||
logger.info(
|
||||
f"VectorStore: top_k_pos={k_pos}, top_k_neg={k_neg if len(neg_similarities) > 0 else 0}, "
|
||||
f"top_k_pos_sim={top_k_pos_sim:.4f}, top_k_neg_sim={top_k_neg_sim:.4f}, "
|
||||
f"diff={diff:.4f}, adaptive_mult={adaptive_multiplier:.2f}, "
|
||||
f"VectorStore: k={self.k}, k_pos={k_pos}, k_neg={k_neg if len(neg_similarities) > 0 else 0}, "
|
||||
f"avg_pos={avg_pos:.4f}, avg_neg={avg_neg:.4f}, "
|
||||
f"diff={diff:.4f}, score_multiplier={self.score_multiplier}, "
|
||||
f"score_neg_pos={score_neg_pos:.4f}, score_pos_only={score_pos_only:.4f}, "
|
||||
f"pos_mean={pos_mean:.4f}±{pos_std:.4f}[{pos_min:.4f}-{pos_max:.4f}], "
|
||||
f"neg_mean={neg_mean:.4f}±{neg_std:.4f}[{neg_min:.4f}-{neg_max:.4f}]"
|
||||
@@ -395,35 +384,88 @@ class VectorStore:
|
||||
|
||||
def _load_from_disk(self) -> None:
|
||||
"""Загружает векторы с диска."""
|
||||
if not self.storage_path or not os.path.exists(self.storage_path):
|
||||
if not self.storage_path:
|
||||
return
|
||||
|
||||
try:
|
||||
with self._lock:
|
||||
data = np.load(self.storage_path, allow_pickle=True)
|
||||
storage_dir = Path(self.storage_path).parent
|
||||
positive_npy = storage_dir / "positive_embeddings.npy"
|
||||
negative_npy = storage_dir / "negative_embeddings.npy"
|
||||
|
||||
# Загружаем векторы
|
||||
pos_vectors = data.get('positive_vectors', np.array([]))
|
||||
neg_vectors = data.get('negative_vectors', np.array([]))
|
||||
# Проверяем наличие отдельных .npy файлов
|
||||
if positive_npy.exists() or negative_npy.exists():
|
||||
logger.info("VectorStore: Обнаружены отдельные .npy файлы, загружаем их...")
|
||||
|
||||
# Загружаем положительные векторы
|
||||
if positive_npy.exists():
|
||||
pos_vectors = np.load(positive_npy, allow_pickle=False)
|
||||
if pos_vectors.size > 0:
|
||||
# Проверяем размерность
|
||||
if len(pos_vectors.shape) == 2:
|
||||
# Массив векторов [N, dim]
|
||||
self._positive_vectors = [vec for vec in pos_vectors]
|
||||
elif len(pos_vectors.shape) == 1:
|
||||
# Один вектор [dim]
|
||||
self._positive_vectors = [pos_vectors]
|
||||
else:
|
||||
logger.warning(f"VectorStore: Неожиданная размерность positive_embeddings.npy: {pos_vectors.shape}")
|
||||
self._positive_vectors = []
|
||||
logger.info(f"VectorStore: Загружено {len(self._positive_vectors)} положительных векторов из {positive_npy}")
|
||||
|
||||
# Загружаем отрицательные векторы
|
||||
if negative_npy.exists():
|
||||
neg_vectors = np.load(negative_npy, allow_pickle=False)
|
||||
if neg_vectors.size > 0:
|
||||
# Проверяем размерность
|
||||
if len(neg_vectors.shape) == 2:
|
||||
# Массив векторов [N, dim]
|
||||
self._negative_vectors = [vec for vec in neg_vectors]
|
||||
elif len(neg_vectors.shape) == 1:
|
||||
# Один вектор [dim]
|
||||
self._negative_vectors = [neg_vectors]
|
||||
else:
|
||||
logger.warning(f"VectorStore: Неожиданная размерность negative_embeddings.npy: {neg_vectors.shape}")
|
||||
self._negative_vectors = []
|
||||
logger.info(f"VectorStore: Загружено {len(self._negative_vectors)} отрицательных векторов из {negative_npy}")
|
||||
|
||||
# Нормализуем загруженные векторы
|
||||
self._positive_vectors = [self._normalize_vector(np.array(v)) for v in self._positive_vectors]
|
||||
self._negative_vectors = [self._normalize_vector(np.array(v)) for v in self._negative_vectors]
|
||||
|
||||
logger.info(
|
||||
f"VectorStore: Загружено с диска из .npy файлов ({self.positive_count} pos, "
|
||||
f"{self.negative_count} neg)"
|
||||
)
|
||||
return
|
||||
|
||||
if pos_vectors.size > 0:
|
||||
self._positive_vectors = list(pos_vectors)
|
||||
if neg_vectors.size > 0:
|
||||
self._negative_vectors = list(neg_vectors)
|
||||
|
||||
# Загружаем хеши
|
||||
pos_hashes = data.get('positive_hashes', np.array([]))
|
||||
neg_hashes = data.get('negative_hashes', np.array([]))
|
||||
|
||||
if pos_hashes.size > 0:
|
||||
self._positive_hashes = list(pos_hashes)
|
||||
if neg_hashes.size > 0:
|
||||
self._negative_hashes = list(neg_hashes)
|
||||
|
||||
logger.info(
|
||||
f"VectorStore: Загружено с диска ({self.positive_count} pos, "
|
||||
f"{self.negative_count} neg): {self.storage_path}"
|
||||
)
|
||||
# Если отдельных .npy файлов нет, пытаемся загрузить из старого формата .npz
|
||||
if os.path.exists(self.storage_path):
|
||||
logger.info(f"VectorStore: Загружаем из старого формата .npz: {self.storage_path}")
|
||||
data = np.load(self.storage_path, allow_pickle=True)
|
||||
|
||||
# Загружаем векторы
|
||||
pos_vectors = data.get('positive_vectors', np.array([]))
|
||||
neg_vectors = data.get('negative_vectors', np.array([]))
|
||||
|
||||
if pos_vectors.size > 0:
|
||||
self._positive_vectors = list(pos_vectors)
|
||||
if neg_vectors.size > 0:
|
||||
self._negative_vectors = list(neg_vectors)
|
||||
|
||||
# Загружаем хеши
|
||||
pos_hashes = data.get('positive_hashes', np.array([]))
|
||||
neg_hashes = data.get('negative_hashes', np.array([]))
|
||||
|
||||
if pos_hashes.size > 0:
|
||||
self._positive_hashes = list(pos_hashes)
|
||||
if neg_hashes.size > 0:
|
||||
self._negative_hashes = list(neg_hashes)
|
||||
|
||||
logger.info(
|
||||
f"VectorStore: Загружено с диска ({self.positive_count} pos, "
|
||||
f"{self.negative_count} neg): {self.storage_path}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"VectorStore: Ошибка загрузки с диска: {e}")
|
||||
@@ -452,26 +494,20 @@ class VectorStore:
|
||||
"""Возвращает текущие параметры формулы расчета score."""
|
||||
return {
|
||||
"score_multiplier": self.score_multiplier,
|
||||
"k_min": self.k_min,
|
||||
"k_max": self.k_max,
|
||||
"base_multiplier_factor": self.base_multiplier_factor,
|
||||
"k": self.k,
|
||||
}
|
||||
|
||||
def update_scoring_params(
|
||||
self,
|
||||
score_multiplier: Optional[float] = None,
|
||||
k_min: Optional[int] = None,
|
||||
k_max: Optional[int] = None,
|
||||
base_multiplier_factor: Optional[float] = None,
|
||||
k: Optional[int] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Обновляет параметры формулы расчета score.
|
||||
|
||||
Args:
|
||||
score_multiplier: Базовый множитель (должен быть > 0)
|
||||
k_min: Минимальное значение k (должно быть >= 1)
|
||||
k_max: Максимальное значение k (должно быть >= k_min)
|
||||
base_multiplier_factor: Множитель для базового score_multiplier (должен быть > 0)
|
||||
score_multiplier: Множитель для масштабирования разницы (должен быть > 0)
|
||||
k: Количество ближайших примеров для расчета среднего (должно быть >= 1)
|
||||
|
||||
Returns:
|
||||
dict: Обновленные параметры
|
||||
@@ -485,30 +521,14 @@ class VectorStore:
|
||||
raise ValueError("score_multiplier должен быть > 0")
|
||||
self.score_multiplier = score_multiplier
|
||||
|
||||
if k_min is not None:
|
||||
if k_min < 1:
|
||||
raise ValueError("k_min должен быть >= 1")
|
||||
if self.k_max < k_min:
|
||||
raise ValueError("k_min не может быть больше k_max")
|
||||
self.k_min = k_min
|
||||
|
||||
if k_max is not None:
|
||||
if k_max < 1:
|
||||
raise ValueError("k_max должен быть >= 1")
|
||||
if k_max < self.k_min:
|
||||
raise ValueError("k_max не может быть меньше k_min")
|
||||
self.k_max = k_max
|
||||
|
||||
if base_multiplier_factor is not None:
|
||||
if base_multiplier_factor <= 0:
|
||||
raise ValueError("base_multiplier_factor должен быть > 0")
|
||||
self.base_multiplier_factor = base_multiplier_factor
|
||||
if k is not None:
|
||||
if k < 1:
|
||||
raise ValueError("k должен быть >= 1")
|
||||
self.k = k
|
||||
|
||||
logger.info(
|
||||
f"VectorStore: Параметры формулы обновлены: "
|
||||
f"score_multiplier={self.score_multiplier}, "
|
||||
f"k_min={self.k_min}, k_max={self.k_max}, "
|
||||
f"base_multiplier_factor={self.base_multiplier_factor}"
|
||||
f"score_multiplier={self.score_multiplier}, k={self.k}"
|
||||
)
|
||||
|
||||
return self.get_scoring_params()
|
||||
|
||||
Reference in New Issue
Block a user