Improve topic sorting: add popular sorting by publications and authors count

services/pretopic.py

@@ -1,170 +1,90 @@
+import asyncio
 import concurrent.futures
-from typing import Dict, List, Tuple
+from concurrent.futures import Future
+from typing import Any, Optional
 
-from txtai.embeddings import Embeddings
+try:
+    from utils.logger import root_logger as logger
+except ImportError:
+    import logging
 
-from services.logger import root_logger as logger
+    logger = logging.getLogger(__name__)
 
 
-class TopicClassifier:
-    def __init__(self, shouts_by_topic: Dict[str, str], publications: List[Dict[str, str]]):
-        """
-        Инициализация классификатора тем и поиска публикаций.
-        Args:
-            shouts_by_topic: Словарь {тема: текст_всех_публикаций}
-            publications: Список публикаций с полями 'id', 'title', 'text'
-        """
-        self.shouts_by_topic = shouts_by_topic
-        self.topics = list(shouts_by_topic.keys())
-        self.publications = publications
-        self.topic_embeddings = None  # Для классификации тем
-        self.search_embeddings = None  # Для поиска публикаций
-        self._initialization_future = None
-        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+class PreTopicService:
+    def __init__(self) -> None:
+        self.topic_embeddings: Optional[Any] = None
+        self.search_embeddings: Optional[Any] = None
+        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
+        self._initialization_future: Optional[Future[None]] = None
 
-    def initialize(self) -> None:
-        """
-        Асинхронная инициализация векторных представлений.
-        """
+    def _ensure_initialization(self) -> None:
+        """Ensure embeddings are initialized"""
         if self._initialization_future is None:
             self._initialization_future = self._executor.submit(self._prepare_embeddings)
-            logger.info("Векторизация текстов начата в фоновом режиме...")
 
     def _prepare_embeddings(self) -> None:
-        """
-        Подготавливает векторные представления для тем и поиска.
-        """
-        logger.info("Начинается подготовка векторных представлений...")
-
-        # Модель для русского языка
-        # TODO: model local caching
-        model_path = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-
-        # Инициализируем embeddings для классификации тем
-        self.topic_embeddings = Embeddings(path=model_path)
-        topic_documents = [(topic, text) for topic, text in self.shouts_by_topic.items()]
-        self.topic_embeddings.index(topic_documents)
-
-        # Инициализируем embeddings для поиска публикаций
-        self.search_embeddings = Embeddings(path=model_path)
-        search_documents = [(str(pub["id"]), f"{pub['title']} {pub['text']}") for pub in self.publications]
-        self.search_embeddings.index(search_documents)
-
-        logger.info("Подготовка векторных представлений завершена.")
-
-    def predict_topic(self, text: str) -> Tuple[float, str]:
-        """
-        Предсказывает тему для заданного текста из известного набора тем.
-        Args:
-            text: Текст для классификации
-        Returns:
-            Tuple[float, str]: (уверенность, тема)
-        """
-        if not self.is_ready():
-            logger.error("Векторные представления не готовы. Вызовите initialize() и дождитесь завершения.")
-            return 0.0, "unknown"
-
+        """Prepare embeddings for topic and search functionality"""
         try:
-            # Ищем наиболее похожую тему
-            results = self.topic_embeddings.search(text, 1)
-            if not results:
-                return 0.0, "unknown"
+            from txtai.embeddings import Embeddings  # type: ignore[import-untyped]
 
-            score, topic = results[0]
-            return float(score), topic
+            # Initialize topic embeddings
+            self.topic_embeddings = Embeddings(
+                {
+                    "method": "transformers",
+                    "path": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+                }
+            )
 
+            # Initialize search embeddings
+            self.search_embeddings = Embeddings(
+                {
+                    "method": "transformers",
+                    "path": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+                }
+            )
+            logger.info("PreTopic embeddings initialized successfully")
+        except ImportError:
+            logger.warning("txtai.embeddings not available, PreTopicService disabled")
         except Exception as e:
-            logger.error(f"Ошибка при определении темы: {str(e)}")
-            return 0.0, "unknown"
+            logger.error(f"Failed to initialize embeddings: {e}")
 
-    def search_similar(self, query: str, limit: int = 5) -> List[Dict[str, any]]:
-        """
-        Ищет публикации похожие на поисковый запрос.
-        Args:
-            query: Поисковый запрос
-            limit: Максимальное количество результатов
-        Returns:
-            List[Dict]: Список найденных публикаций с оценкой релевантности
-        """
-        if not self.is_ready():
-            logger.error("Векторные представления не готовы. Вызовите initialize() и дождитесь завершения.")
+    async def suggest_topics(self, text: str) -> list[dict[str, Any]]:
+        """Suggest topics based on text content"""
+        if self.topic_embeddings is None:
             return []
 
         try:
-            # Ищем похожие публикации
-            results = self.search_embeddings.search(query, limit)
-
-            # Формируем результаты
-            found_publications = []
-            for score, pub_id in results:
-                # Находим публикацию по id
-                publication = next((pub for pub in self.publications if str(pub["id"]) == pub_id), None)
-                if publication:
-                    found_publications.append({**publication, "relevance": float(score)})
-
-            return found_publications
+            self._ensure_initialization()
+            if self._initialization_future:
+                await asyncio.wrap_future(self._initialization_future)
 
+            if self.topic_embeddings is not None:
+                results = self.topic_embeddings.search(text, 1)
+                if results:
+                    return [{"topic": result["text"], "score": result["score"]} for result in results]
         except Exception as e:
-            logger.error(f"Ошибка при поиске публикаций: {str(e)}")
+            logger.error(f"Error suggesting topics: {e}")
+        return []
+
+    async def search_content(self, query: str, limit: int = 10) -> list[dict[str, Any]]:
+        """Search content using embeddings"""
+        if self.search_embeddings is None:
             return []
 
-    def is_ready(self) -> bool:
-        """
-        Проверяет, готовы ли векторные представления.
-        """
-        return self.topic_embeddings is not None and self.search_embeddings is not None
+        try:
+            self._ensure_initialization()
+            if self._initialization_future:
+                await asyncio.wrap_future(self._initialization_future)
 
-    def wait_until_ready(self) -> None:
-        """
-        Ожидает завершения подготовки векторных представлений.
-        """
-        if self._initialization_future:
-            self._initialization_future.result()
-
-    def __del__(self):
-        """
-        Очистка ресурсов при удалении объекта.
-        """
-        if self._executor:
-            self._executor.shutdown(wait=False)
+            if self.search_embeddings is not None:
+                results = self.search_embeddings.search(query, limit)
+                if results:
+                    return [{"content": result["text"], "score": result["score"]} for result in results]
+        except Exception as e:
+            logger.error(f"Error searching content: {e}")
+        return []
 
 
-# Пример использования:
-"""
-shouts_by_topic = {
-    "Спорт": "... большой текст со всеми спортивными публикациями ...",
-    "Технологии": "... большой текст со всеми технологическими публикациями ...",
-    "Политика": "... большой текст со всеми политическими публикациями ..."
-}
-
-publications = [
-    {
-        'id': 1,
-        'title': 'Новый процессор AMD',
-        'text': 'Компания AMD представила новый процессор...'
-    },
-    {
-        'id': 2,
-        'title': 'Футбольный матч',
-        'text': 'Вчера состоялся решающий матч...'
-    }
-]
-
-# Создание классификатора
-classifier = TopicClassifier(shouts_by_topic, publications)
-classifier.initialize()
-classifier.wait_until_ready()
-
-# Определение темы текста
-text = "Новый процессор показал высокую производительность"
-score, topic = classifier.predict_topic(text)
-print(f"Тема: {topic} (уверенность: {score:.4f})")
-
-# Поиск похожих публикаций
-query = "процессор AMD производительность"
-similar_publications = classifier.search_similar(query, limit=3)
-for pub in similar_publications:
-    print(f"\nНайдена публикация (релевантность: {pub['relevance']:.4f}):")
-    print(f"Заголовок: {pub['title']}")
-    print(f"Текст: {pub['text'][:100]}...")
-"""
+# Global instance
+pretopic_service = PreTopicService()