OptimezedSumAndUserSearch

app.py

@@ -173,7 +173,7 @@ def get_whisper_pipeline() -> Any:
             model=model,
             tokenizer=_WHISPER_PROCESSOR.tokenizer,
             feature_extractor=_WHISPER_PROCESSOR.feature_extractor,
-            torch_dtype=torch.float32,
+            dtype=torch.float32,
             device="cpu",
         )
 

extractors/amount_extractor.py

@@ -1,60 +1,29 @@
-"""
-Экстрактор сумм из текста.
-
-Использует GLiNER для извлечения денежных сумм.
-"""
+"""Простой regex-экстрактор суммы из текста."""
 
 from __future__ import annotations
 
 import re
 from typing import Any, Optional
 
-from gliner import GLiNER
-
-
-# Глобальная модель для извлечения сумм
-_AMOUNT_MODEL: Optional[GLiNER] = None
-
 
-def get_amount_model() -> Optional[GLiNER]:
-    """Возвращает модель для извлечения сумм (ленивая загрузка)."""
-    global _AMOUNT_MODEL
-
-    if _AMOUNT_MODEL is None:
-        _AMOUNT_MODEL = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
-
-    return _AMOUNT_MODEL
+AMOUNT_PATTERN = re.compile(r"\d+(?:,\d{1,2})?", re.IGNORECASE)
 
 
 class ExpenseAmountExtractor:
-    """
-    Экстрактор денежных сумм из текста.
-    
-    Использует GLiNER для поиска упоминаний денег.
-    """
-    
+    """Извлекает сумму как целое число или число с запятой."""
+
     def __init__(self, suppliers: list[str] | None = None) -> None:
-        """
-        Args:
-            suppliers: Список поставщиков (не используется, для совместимости)
-        """
-        self.model = get_amount_model()
+        self.suppliers = suppliers or []
 
     @staticmethod
     def to_float(value: str) -> Optional[float]:
-        """Преобразует строку в число."""
-        cleaned = value.replace(" ", "").replace("\u00A0", "")
-        match = re.search(r"\d+(?:[,]\d{1,2})?", cleaned)
-        if not match:
-            return None
         try:
-            return float(match.group(0).replace(",", "."))
+            return float(value.replace(",", "."))
         except ValueError:
             return None
 
     @staticmethod
     def phrase_span(text: str, phrase: Optional[str]) -> Optional[tuple[int, int]]:
-        """Возвращает позицию фразы в тексте."""
         if not phrase:
             return None
         idx = text.lower().find(phrase.lower())
@@ -64,61 +33,30 @@ class ExpenseAmountExtractor:
 
     @staticmethod
     def overlaps(span1: tuple[int, int], span2: Optional[tuple[int, int]]) -> bool:
-        """Проверяет пересечение двух диапазонов."""
         if span2 is None:
             return False
         return span1[0] < span2[1] and span2[0] < span1[1]
 
-    @staticmethod
-    def expand_amount_text(text: str, start: int, end: int) -> tuple[str, tuple[int, int]]:
-        """Расширяет текст суммы (для дробных чисел)."""
-        suffix = re.match(r",\d{1,2}", text[end:])
-        if suffix:
-            new_end = end + len(suffix.group(0))
-            return text[start:new_end].strip(), (start, new_end)
-
-        prefix = re.search(r"(\d{1,3}(?:\s*\d{3})*),", text[:start])
-        if prefix:
-            new_start = prefix.start(1)
-            return text[new_start:end].strip(), (new_start, end)
-
-        return text[start:end].strip(), (start, end)
-
     def extract(
         self,
         text: str,
         matched_date_phrase: Optional[str] = None,
         matched_supplier_phrase: Optional[str] = None,
     ) -> dict[str, Any]:
-        """
-        Извлекает сумму из текста.
-        
-        Args:
-            text: Текст для анализа
-            matched_date_phrase: Фраза даты для исключения
-            matched_supplier_phrase: Фраза поставщика для исключения
-            
-        Returns:
-            Словарь с amount и amount_text
-        """
-        if self.model is None:
-            return {"amount": None, "amount_text": None}
-
         date_span = self.phrase_span(text, matched_date_phrase)
         supplier_span = self.phrase_span(text, matched_supplier_phrase)
-        entities = self.model.predict_entities(text, ["money"], threshold=0.3)
 
-        for ent in sorted(entities, key=lambda item: float(item.get("score", 0.0)), reverse=True):
-            raw_span = (int(ent.get("start", 0)), int(ent.get("end", 0)))
-            amount_text, span = self.expand_amount_text(text, raw_span[0], raw_span[1])
-            amount = self.to_float(amount_text)
-            overlaps_date = self.overlaps(span, date_span)
-            overlaps_supplier = self.overlaps(span, supplier_span)
+        for match in AMOUNT_PATTERN.finditer(text):
+            span = match.span()
 
-            if amount is None:
+            if self.overlaps(span, date_span):
                 continue
-            if overlaps_date or overlaps_supplier:
+            if self.overlaps(span, supplier_span):
                 continue
-            return {"amount": amount, "amount_text": amount_text}
+
+            amount_text = match.group(0)
+            amount = self.to_float(amount_text)
+            if amount is not None:
+                return {"amount": amount, "amount_text": amount_text}
 
         return {"amount": None, "amount_text": None}

extractors/user_extractor.py

@@ -43,98 +43,109 @@ def lemmatize_text(text: str) -> list[str]:
 class ExpenseUserExtractor:
     """
     Экстрактор пользователей из текста.
-    
-    Использует семантические эмбеддинги для сопоставления слов из текста
-    с известными пользователями.
+
+    Сначала использует точное/текстовое совпадение имени, а затем
+    один батч эмбеддингов для оставшихся кандидатов.
     """
-    
+
     def __init__(
-        self, 
-        users: list[str], 
-        suppliers: list[str], 
-        model: SentenceTransformer, 
+        self,
+        users: list[str],
+        suppliers: list[str],
+        model: SentenceTransformer,
         threshold: float = 0.6
     ) -> None:
-        """
-        Args:
-            users: Список известных пользователей
-            suppliers: Список поставщиков (для исключения)
-            model: Модель для создания эмбеддингов
-            threshold: Порог схожести
-        """
         self.users = users
         self.model = model
         self.threshold = threshold
         self.supplier_terms = {normalize_text(supplier) for supplier in suppliers}
         self.user_terms = [normalize_text(user) for user in users]
+        self.user_lookup = dict(zip(self.user_terms, self.users))
         self.user_embeddings = model.encode(
             [f"passage: {user}" for user in self.user_terms],
             convert_to_tensor=True,
             normalize_embeddings=True,
+            show_progress_bar=False,
         )
 
+    @staticmethod
+    def _contains_whole_phrase(text: str, phrase: str) -> bool:
+        if not phrase:
+            return False
+        return re.search(rf"(?<!\w){re.escape(phrase)}(?!\w)", text) is not None
+
+    def _extract_candidates(self, text: str, excluded_tokens: set[str]) -> list[str]:
+        candidates: list[str] = []
+        seen: set[str] = set()
+
+        for word in lemmatize_text(text):
+            if len(word) < 3:
+                continue
+            if word in excluded_tokens or word in self.supplier_terms or word in seen:
+                continue
+            seen.add(word)
+            candidates.append(word)
+
+        return candidates
+
     def extract(
-        self, 
-        text: str, 
-        supplier_phrase: str | None = None, 
+        self,
+        text: str,
+        supplier_phrase: str | None = None,
         date_phrase: str | None = None
     ) -> dict[str, Any]:
-        """
-        Извлекает пользователя из текста.
-        
-        Args:
-            text: Текст для анализа
-            supplier_phrase: Фраза поставщика для исключения
-            date_phrase: Фраза даты для исключения
-            
-        Returns:
-            Словарь с user, user_score, matched_user_phrase
-        """
         excluded_tokens: set[str] = set()
         if supplier_phrase:
             excluded_tokens.update(normalize_text(supplier_phrase).split())
         if date_phrase:
             excluded_tokens.update(normalize_text(date_phrase).split())
 
-        best_user = None
-        best_score = -1.0
-        best_phrase = None
-
-        for word in lemmatize_text(text):
-            if len(word) < 3:
-                continue
-            if word in excluded_tokens or word in self.supplier_terms:
-                continue
+        normalized_text = normalize_text(text)
 
-            query_emb = self.model.encode(
-                f"query: {word}",
-                convert_to_tensor=True,
-                normalize_embeddings=True,
-            )
-            similarities = torch.cosine_similarity(query_emb.unsqueeze(0), self.user_embeddings, dim=1)
-            idx = int(torch.argmax(similarities))
-            score = similarities[idx].item()
-
-            if score > best_score:
-                best_score = score
-                best_user = self.users[idx]
-                best_phrase = word
-
-        if best_score >= self.threshold:
-            return {
-                "user": best_user,
-                "user_score": round(best_score, 4),
-                "matched_user_phrase": best_phrase,
-            }
+        for user_term, original_user in zip(self.user_terms, self.users):
+            if user_term and self._contains_whole_phrase(normalized_text, user_term):
+                return {
+                    "user": original_user,
+                    "user_score": 1.0,
+                    "matched_user_phrase": user_term,
+                }
 
-        # Проверка на местоимение "я"
-        if re.search(r"(?<!\S)я(?!\S)", normalize_text(text), re.IGNORECASE):
+        if re.search(r"(?<!\S)я(?!\S)", normalized_text, re.IGNORECASE):
             return {
                 "user": "Я",
                 "user_score": 1.0,
                 "matched_user_phrase": "я",
             }
 
+        candidates = self._extract_candidates(text, excluded_tokens)
+        if not candidates:
+            return {
+                "user": None,
+                "user_score": None,
+                "matched_user_phrase": None,
+            }
+
+        query_embeddings = self.model.encode(
+            [f"query: {word}" for word in candidates],
+            convert_to_tensor=True,
+            normalize_embeddings=True,
+            show_progress_bar=False,
+            batch_size=max(1, min(32, len(candidates))),
+        )
+
+        similarity_matrix = torch.matmul(query_embeddings, self.user_embeddings.T)
+        flat_index = int(torch.argmax(similarity_matrix))
+        candidate_index = flat_index // len(self.users)
+        user_index = flat_index % len(self.users)
+        score = similarity_matrix[candidate_index, user_index].item()
+
+        if score >= self.threshold:
+            return {
+                "user": self.users[user_index],
+                "user_score": round(score, 4),
+                "matched_user_phrase": candidates[candidate_index],
+            }
+
         return {
             "user": None,
             "user_score": None,