diff --git a/Dockerfile b/Dockerfile index 49df5b8..2fd5c6d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -25,7 +25,7 @@ COPY app/ ./app/ RUN mkdir -p data/models data/vectors # Переменные окружения по умолчанию -ENV RAG_MODEL=DeepPavlov/rubert-base-cased +ENV RAG_MODEL=sentence-transformers/all-MiniLM-L12-v2 ENV RAG_CACHE_DIR=/app/data/models ENV RAG_VECTORS_PATH=/app/data/vectors/vectors.npz ENV RAG_API_HOST=0.0.0.0 diff --git a/app/api/routes.py b/app/api/routes.py index 7d2c714..9a45185 100644 --- a/app/api/routes.py +++ b/app/api/routes.py @@ -416,9 +416,7 @@ async def update_scoring_params( try: params = service.vector_store.update_scoring_params( score_multiplier=request.score_multiplier, - k_min=request.k_min, - k_max=request.k_max, - base_multiplier_factor=request.base_multiplier_factor, + k=request.k, ) return ScoringParamsResponse(**params) except ValueError as e: diff --git a/app/config.py b/app/config.py index c98343a..b456e68 100644 --- a/app/config.py +++ b/app/config.py @@ -18,7 +18,7 @@ class Settings: # Модель model_name: str = field( - default_factory=lambda: os.getenv("RAG_MODEL", "DeepPavlov/rubert-base-cased") + default_factory=lambda: os.getenv("RAG_MODEL", "sentence-transformers/all-MiniLM-L12-v2") ) cache_dir: str = field( default_factory=lambda: os.getenv("RAG_CACHE_DIR", "data/models") @@ -73,8 +73,8 @@ class Settings: default_factory=lambda: int(os.getenv("RAG_AUTOSAVE_INTERVAL", "600")) # 10 минут ) - # Размерность векторов (768 для ruBERT) - vector_dim: int = 768 + # Размерность векторов (384 для all-MiniLM-L12-v2) + vector_dim: int = 384 @property def is_auth_required(self) -> bool: diff --git a/app/main.py b/app/main.py index 45a08e3..f871419 100644 --- a/app/main.py +++ b/app/main.py @@ -1,7 +1,7 @@ """ -FastAPI приложение RAG сервиса. +FastAPI приложение Embedding сервиса. -Сервис для векторного скоринга текстов с использованием ruBERT. +Сервис для векторного скоринга текстов с использованием sentence-transformers. """ import asyncio @@ -90,7 +90,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: global _autosave_task setup_logging() - logger.info(f"RAG Service v{__version__} запускается...") + logger.info(f"Embedding Service v{__version__} запускается...") settings = get_settings() logger.info(f"Настройки: model={settings.model_name}, vectors_path={settings.vectors_path}") @@ -112,7 +112,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # logger.info("Прогрев модели при запуске...") # await service.warmup() - logger.info("RAG Service готов к работе") + logger.info("Embedding Service готов к работе") yield @@ -125,21 +125,21 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: pass # При остановке сохраняем векторы - logger.info("RAG Service останавливается, финальное сохранение векторов...") + logger.info("Embedding Service останавливается, финальное сохранение векторов...") try: service.save_vectors() logger.info("Векторы сохранены") except Exception as e: logger.error(f"Ошибка сохранения векторов: {e}") - logger.info("RAG Service остановлен") + logger.info("Embedding Service остановлен") # Создание приложения app = FastAPI( - title="RAG Service", + title="Embedding Service", description=""" - Сервис векторного скоринга текстов с использованием ruBERT. + Сервис векторного скоринга текстов с использованием sentence-transformers. ## Возможности @@ -150,10 +150,10 @@ app = FastAPI( ## Алгоритм скоринга - 1. Текст преобразуется в вектор через ruBERT (768 измерений) - 2. Вычисляется косинусное сходство с положительными примерами - 3. Вычисляется косинусное сходство с отрицательными примерами - 4. Финальный скор = разница между сходствами, нормализованная в [0, 1] + 1. Текст преобразуется в вектор через sentence-transformers/all-MiniLM-L12-v2 (384 измерения) + 2. Вычисляется косинусное сходство с положительными примерами (топ-k ближайших) + 3. Вычисляется косинусное сходство с отрицательными примерами (топ-k ближайших) + 4. Финальный скор = (diff * multiplier + 1) / 2, где diff = avg_pos - avg_neg, нормализованный в [0, 1] """, version=__version__, lifespan=lifespan, diff --git a/app/schemas.py b/app/schemas.py index 69c86d4..77c2485 100644 --- a/app/schemas.py +++ b/app/schemas.py @@ -1,5 +1,5 @@ """ -Pydantic схемы для API RAG сервиса. +Pydantic схемы для API Embedding сервиса. """ from typing import Any, Dict, Optional @@ -65,7 +65,7 @@ class ScoreResponse(BaseModel): "meta": { "positive_examples": 500, "negative_examples": 350, - "model": "DeepPavlov/rubert-base-cased", + "model": "sentence-transformers/all-MiniLM-L12-v2", "timestamp": 1706270000 } } @@ -111,14 +111,14 @@ class StatsResponse(BaseModel): model_config = { "json_schema_extra": { "example": { - "model_name": "DeepPavlov/rubert-base-cased", + "model_name": "sentence-transformers/all-MiniLM-L12-v2", "model_loaded": True, "device": "cpu", "vector_store": { "positive_count": 500, "negative_count": 350, "total_count": 850, - "vector_dim": 768, + "vector_dim": 384, "max_examples": 10000 } } @@ -180,39 +180,21 @@ class ScoringParamsResponse(BaseModel): score_multiplier: float = Field( ..., description=( - "Базовый множитель для усиления разницы в скорах. " - "Используется как основа для расчета финального множителя. " + "Множитель для масштабирования разницы в скорах. " + "Используется в формуле: score = (diff * score_multiplier + 1) / 2, " + "где diff = avg_pos - avg_neg (разница средних сходств топ-k примеров). " "Чем больше значение, тем сильнее влияние разницы между положительными и отрицательными примерами на итоговый score. " "Рекомендуемое значение: 5.0" ) ) - k_min: int = Field( + k: int = Field( ..., description=( - "Минимальное количество ближайших примеров для расчета среднего сходства. " + "Количество ближайших примеров для расчета среднего сходства. " "Алгоритм берет топ-k самых похожих примеров из каждого типа (положительные/отрицательные) " "и вычисляет среднее косинусное сходство. " "Меньшее значение k делает алгоритм более чувствительным к различиям, но может быть менее стабильным. " - "Рекомендуемое значение: 5" - ) - ) - k_max: int = Field( - ..., - description=( - "Максимальное количество ближайших примеров для расчета среднего сходства. " - "Алгоритм выбирает k в диапазоне [k_min, k_max] в зависимости от количества доступных примеров. " - "Большее значение k делает алгоритм более стабильным, но менее чувствительным к различиям. " - "Должно быть >= k_min. Рекомендуемое значение: 10" - ) - ) - base_multiplier_factor: float = Field( - ..., - description=( - "Множитель для базового score_multiplier. " - "Финальный множитель рассчитывается как: score_multiplier * base_multiplier_factor * адаптивный_коэффициент. " - "Этот параметр усиливает влияние разницы между положительными и отрицательными примерами. " - "Чем больше значение, тем больше диапазон итогового score (от 0 до 1). " - "Рекомендуемое значение: 15.0" + "Рекомендуемое значение: 3" ) ) @@ -220,9 +202,7 @@ class ScoringParamsResponse(BaseModel): "json_schema_extra": { "example": { "score_multiplier": 5.0, - "k_min": 5, - "k_max": 10, - "base_multiplier_factor": 15.0 + "k": 3 } } } @@ -234,42 +214,22 @@ class UpdateScoringParamsRequest(BaseModel): None, gt=0, description=( - "Базовый множитель для усиления разницы в скорах. " - "Используется как основа для расчета финального множителя. " + "Множитель для масштабирования разницы в скорах. " + "Используется в формуле: score = (diff * score_multiplier + 1) / 2, " + "где diff = avg_pos - avg_neg (разница средних сходств топ-k примеров). " "Чем больше значение, тем сильнее влияние разницы между положительными и отрицательными примерами на итоговый score. " "Должен быть > 0. Рекомендуемое значение: 5.0" ) ) - k_min: Optional[int] = Field( + k: Optional[int] = Field( None, ge=1, description=( - "Минимальное количество ближайших примеров для расчета среднего сходства. " + "Количество ближайших примеров для расчета среднего сходства. " "Алгоритм берет топ-k самых похожих примеров из каждого типа (положительные/отрицательные) " "и вычисляет среднее косинусное сходство. " "Меньшее значение k делает алгоритм более чувствительным к различиям, но может быть менее стабильным. " - "Должно быть >= 1. Рекомендуемое значение: 5" - ) - ) - k_max: Optional[int] = Field( - None, - ge=1, - description=( - "Максимальное количество ближайших примеров для расчета среднего сходства. " - "Алгоритм выбирает k в диапазоне [k_min, k_max] в зависимости от количества доступных примеров. " - "Большее значение k делает алгоритм более стабильным, но менее чувствительным к различиям. " - "Должно быть >= 1 и >= k_min. Рекомендуемое значение: 10" - ) - ) - base_multiplier_factor: Optional[float] = Field( - None, - gt=0, - description=( - "Множитель для базового score_multiplier. " - "Финальный множитель рассчитывается как: score_multiplier * base_multiplier_factor * адаптивный_коэффициент. " - "Этот параметр усиливает влияние разницы между положительными и отрицательными примерами. " - "Чем больше значение, тем больше диапазон итогового score (от 0 до 1). " - "Должен быть > 0. Рекомендуемое значение: 15.0" + "Должно быть >= 1. Рекомендуемое значение: 3" ) ) @@ -277,9 +237,7 @@ class UpdateScoringParamsRequest(BaseModel): "json_schema_extra": { "example": { "score_multiplier": 5.0, - "k_min": 5, - "k_max": 10, - "base_multiplier_factor": 15.0 + "k": 3 } } } diff --git a/app/services/rag_service.py b/app/services/rag_service.py index 990b95b..204cf17 100644 --- a/app/services/rag_service.py +++ b/app/services/rag_service.py @@ -1,7 +1,7 @@ """ -RAG сервис для скоринга постов с использованием ruBERT. +RAG сервис для скоринга постов с использованием sentence-transformers. -Использует модель DeepPavlov/rubert-base-cased для создания эмбеддингов +Использует модель sentence-transformers/all-MiniLM-L12-v2 для создания эмбеддингов и сравнивает их с эталонными примерами через VectorStore. """ @@ -66,7 +66,7 @@ class RAGService: """ RAG сервис для оценки постов на основе векторного сходства. - Использует ruBERT для создания эмбеддингов текста и сравнивает + Использует sentence-transformers для создания эмбеддингов текста и сравнивает их с эталонными примерами (опубликованные vs отклоненные посты). Attributes: @@ -92,9 +92,8 @@ class RAGService: self.cache_dir = self._settings.cache_dir self.min_text_length = self._settings.min_text_length - # Модель и токенизатор загружаются лениво + # Модель загружается лениво self._model = None - self._tokenizer = None self._device = None self._model_loaded = False @@ -104,6 +103,7 @@ class RAGService: max_examples=self._settings.max_examples, storage_path=self._settings.vectors_path, score_multiplier=self._settings.score_multiplier, + k=3, # Фиксированное значение k для топ-k ближайших примеров ) logger.info(f"RAGService инициализирован (model={self.model_name})") @@ -138,64 +138,37 @@ class RAGService: def _load_model_sync(self) -> None: """Синхронная загрузка модели (вызывается в executor).""" - logger.info("RAGService: Начало _load_model_sync, импорт transformers...") - from transformers import AutoModel, AutoTokenizer + logger.info("RAGService: Начало _load_model_sync, импорт sentence_transformers...") + from sentence_transformers import SentenceTransformer import torch # Определяем устройство self._device = "cuda" if torch.cuda.is_available() else "cpu" logger.info(f"RAGService: Устройство определено: {self._device}") - # Загружаем токенизатор - logger.info(f"RAGService: Загрузка токенизатора из {self.model_name}...") - self._tokenizer = AutoTokenizer.from_pretrained( - self.model_name, - cache_dir=self.cache_dir, - ) - logger.info("RAGService: Токенизатор загружен") - - # Загружаем модель + # Загружаем модель SentenceTransformer logger.info(f"RAGService: Загрузка модели из {self.model_name} (это может занять несколько минут)...") - self._model = AutoModel.from_pretrained( + self._model = SentenceTransformer( self.model_name, - cache_dir=self.cache_dir, + cache_folder=self.cache_dir, + device=self._device, ) - logger.info("RAGService: Модель загружена, перенос на устройство...") - self._model.to(self._device) - self._model.eval() # Режим инференса - logger.info(f"RAGService: Модель готова на устройстве: {self._device}") def _get_embedding_sync(self, text: str) -> np.ndarray: """ Получает эмбеддинг текста (синхронно). - Использует [CLS] токен как представление всего текста. + Использует SentenceTransformer для получения нормализованного эмбеддинга. Args: text: Текст для векторизации Returns: - Numpy массив с эмбеддингом (768 измерений для ruBERT) + Numpy массив с эмбеддингом (384 измерений для all-MiniLM-L12-v2) """ - import torch - - # Токенизация с ограничением длины - inputs = self._tokenizer( - text, - return_tensors="pt", - truncation=True, - max_length=512, - padding=True, - ) - inputs = {k: v.to(self._device) for k, v in inputs.items()} - - # Получаем эмбеддинг - with torch.no_grad(): - outputs = self._model(**inputs) - # Используем [CLS] токен (первый токен) - embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy() - + # SentenceTransformer автоматически нормализует эмбеддинги + embedding = self._model.encode(text, convert_to_numpy=True, normalize_embeddings=True) return embedding.flatten() def _get_embeddings_batch_sync(self, texts: List[str], batch_size: int = 16) -> List[np.ndarray]: @@ -211,37 +184,17 @@ class RAGService: Returns: Список numpy массивов с эмбеддингами """ - import torch + # SentenceTransformer автоматически обрабатывает батчи + embeddings = self._model.encode( + texts, + batch_size=batch_size, + convert_to_numpy=True, + normalize_embeddings=True, + show_progress_bar=False, + ) - all_embeddings = [] - - for i in range(0, len(texts), batch_size): - batch_texts = texts[i:i + batch_size] - - # Токенизация батча - inputs = self._tokenizer( - batch_texts, - return_tensors="pt", - truncation=True, - max_length=512, - padding=True, - ) - inputs = {k: v.to(self._device) for k, v in inputs.items()} - - # Получаем эмбеддинги - with torch.no_grad(): - outputs = self._model(**inputs) - # [CLS] токен для каждого текста в батче - batch_embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy() - - # Разбиваем на отдельные эмбеддинги - for j in range(len(batch_texts)): - all_embeddings.append(batch_embeddings[j]) - - if i > 0 and i % (batch_size * 10) == 0: - logger.info(f"RAGService: Обработано {i}/{len(texts)} текстов") - - return all_embeddings + # Преобразуем в список отдельных массивов + return [emb.flatten() for emb in embeddings] async def get_embeddings_batch(self, texts: List[str], batch_size: Optional[int] = None) -> List[np.ndarray]: """ diff --git a/app/storage/vector_store.py b/app/storage/vector_store.py index 4eaec69..436b1fe 100644 --- a/app/storage/vector_store.py +++ b/app/storage/vector_store.py @@ -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() diff --git a/docker-compose.yml b/docker-compose.yml index 64b28cf..4539989 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -16,7 +16,7 @@ services: - ./data/models:/app/data/models - ./data/vectors:/app/data/vectors environment: - - RAG_MODEL=${RAG_MODEL:-DeepPavlov/rubert-base-cased} + - RAG_MODEL=${RAG_MODEL:-sentence-transformers/all-MiniLM-L12-v2} - RAG_CACHE_DIR=/app/data/models - RAG_VECTORS_PATH=/app/data/vectors/vectors.npz - RAG_MAX_EXAMPLES=${RAG_MAX_EXAMPLES:-10000} diff --git a/env.example b/env.example index e9f6538..0e289e0 100644 --- a/env.example +++ b/env.example @@ -1,7 +1,7 @@ # RAG Service Configuration # Модель -RAG_MODEL=DeepPavlov/rubert-base-cased +RAG_MODEL=sentence-transformers/all-MiniLM-L12-v2 RAG_CACHE_DIR=data/models # VectorStore diff --git a/pyproject.toml b/pyproject.toml index 9fd66c6..e9434c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rag-service" version = "0.1.0" -description = "RAG Service - сервис векторного скоринга на FastAPI с ruBERT" +description = "Embedding Service - сервис векторного скоринга на FastAPI с sentence-transformers" readme = "README.md" requires-python = ">=3.11" license = {text = "MIT"} diff --git a/requirements.txt b/requirements.txt index 813f854..d28ca2f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,6 +6,7 @@ pydantic>=2.5.0 # ML / NLP torch>=2.1.0 transformers>=4.36.0 +sentence-transformers>=2.2.0 numpy>=1.24.0 # Утилиты