добавлення функціонала для підключення моделей для локального ембедінга

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from sdc_classifier import SDCClassifier
 from dotenv import load_dotenv
+import torch
 import json
 import os
 
@@ -35,6 +36,24 @@ def initialize_environment():
             
     return True
 
+def create_classifier(model_type, openai_model=None, local_model=None, device=None):
+    """
+    Створення класифікатора з відповідними параметрами
+    
+    Args:
+        model_type: тип моделі ("OpenAI" або "Local")
+        openai_model: назва моделі OpenAI
+        local_model: шлях до локальної моделі
+        device: пристрій для локальної моделі
+    
+    Returns:
+        SDCClassifier: налаштований класифікатор
+    """
+    if model_type == "OpenAI":
+        return SDCClassifier()
+    else:
+        return SDCClassifier(local_model=local_model, device=device)
+
 def main():
     # Константи файлів
     DEFAULT_CLASSES_FILE = "classes.json"
@@ -58,7 +77,7 @@ def main():
         try:
             classifier.load_initial_state(DEFAULT_CLASSES_FILE, DEFAULT_SIGNATURES_FILE)
             result = classifier.initialize_signatures(
-                force_rebuild=True,  # Примусово будуємо нові signatures
+                force_rebuild=True,
                 signatures_file=DEFAULT_SIGNATURES_FILE
             )
             print(f"Результат ініціалізації: {result}")
@@ -101,11 +120,33 @@ def main():
                 # Налаштування моделі
                 with gr.Accordion("Налаштування моделі", open=False):
                     with gr.Row():
+                        model_type = gr.Radio(
+                            choices=["OpenAI", "Local"],
+                            value="OpenAI",
+                            label="Тип моделі"
+                        )
                         model_choice = gr.Dropdown(
-                            choices=["text-embedding-3-large","text-embedding-3-small"],
+                            choices=[
+                                "text-embedding-3-large",
+                                "text-embedding-3-small"
+                            ],
                             value="text-embedding-3-large",
-                            label="OpenAI model"
+                            label="OpenAI model",
+                            visible=True
+                        )
+                        local_model_path = gr.Textbox(
+                            value="cambridgeltl/SapBERT-from-PubMedBERT-fulltext",
+                            label="Шлях до локальної моделі",
+                            visible=False
+                        )
+                        device_choice = gr.Radio(
+                            choices=["cuda", "cpu"],
+                            value="cuda" if torch.cuda.is_available() else "cpu",
+                            label="Пристрій для локальної моделі",
+                            visible=False
                         )
+                    
+                    with gr.Row():
                         json_file = gr.File(
                             label="Завантажити новий JSON з класами",
                             file_types=[".json"]
@@ -114,6 +155,7 @@ def main():
                             label="Примусово перебудувати signatures",
                             value=False
                         )
+                        
                     with gr.Row():
                         build_btn = gr.Button("Оновити signatures")
                         build_out = gr.Label(label="Статус signatures")
@@ -156,80 +198,129 @@ def main():
         gr.Markdown("""
         ### Інструкція:
         1. У вкладці "Налаштування моделі" можна:
+           - Вибрати тип моделі (OpenAI або Local)
+           - Налаштувати параметри вибраної моделі
            - Завантажити новий JSON файл з класами
-           - Вибрати модель для embeddings
            - Примусово перебудувати signatures
-        2. Після зміни класів натисніть "Оновити signatures"
+        2. Після зміни налаштувань натисніть "Оновити signatures"
         3. Використовуйте повзунок "Поріг впевненості" для фільтрації результатів
         4. На вкладці "Пакетна обробка" можна аналізувати багато повідомлень
         5. Результати можна зберегти в CSV файл
         """)
         
         # Підключення обробників подій
-        def update_with_file(file, model_name, force):
-            if file is None:
-                # Відновлюємо базовий стан якщо файл видалено
-                classifier.restore_base_state()
-                return ("Відновлено базовий набір класів", classifier.get_cache_stats())
-                
+        def update_model_inputs(model_type):
+            """Оновлення видимості полів в залежності від типу моделі"""
+            return {
+                model_choice: gr.update(visible=model_type == "OpenAI"),
+                local_model_path: gr.update(visible=model_type == "Local"),
+                device_choice: gr.update(visible=model_type == "Local")
+            }
+        
+        def update_classifier_settings(json_file, model_type, openai_model, 
+                                    local_model, device, force_rebuild):
+            """Оновлення налаштувань класифікатора"""
             try:
-                # Для роботи з gradio File компонентом
-                if hasattr(file, 'name'):  # Якщо це файловий об'єкт
-                    with open(file.name, 'r', encoding='utf-8') as f:
-                        new_classes = json.load(f)
-                else:  # Якщо це строка
-                    new_classes = json.loads(file)
-                
-                if not isinstance(new_classes, dict):
-                    return ("Помилка: JSON повинен містити словник класів", classifier.get_cache_stats())
+                # Створюємо новий класифікатор з вибраними параметрами
+                nonlocal classifier
+                classifier = create_classifier(
+                    model_type=model_type,
+                    openai_model=openai_model if model_type == "OpenAI" else None,
+                    local_model=local_model if model_type == "Local" else None,
+                    device=device if model_type == "Local" else None
+                )
                 
-                # Завантажуємо нові класи без перезапису файлу за замовчуванням
-                classifier.load_classes(new_classes)
+                # Завантажуємо класи
+                if json_file is not None:
+                    with open(json_file.name, 'r', encoding='utf-8') as f:
+                        new_classes = json.load(f)
+                    classifier.load_classes(new_classes)
+                else:
+                    classifier.restore_base_state()
                 
-                # Створюємо тимчасові signatures
+                # Ініціалізуємо signatures
                 result = classifier.initialize_signatures(
-                    model_name=model_name,
-                    signatures_file=None,  # Не зберігаємо у файл
-                    force_rebuild=True  # Завжди перебудовуємо для нових класів
+                    force_rebuild=force_rebuild,
+                    signatures_file=DEFAULT_SIGNATURES_FILE if not force_rebuild else None
                 )
                 
-                return (f"Тимчасові класи завантажено. {result}", classifier.get_cache_stats())
-                
-            except json.JSONDecodeError:
-                return ("Помилка: Неправильний формат JSON файлу", classifier.get_cache_stats())
+                return result, classifier.get_cache_stats()
             except Exception as e:
-                return (f"Помилка при оновленні: {str(e)}", classifier.get_cache_stats())
-
-        single_process_btn.click(
-            fn=lambda text, threshold: classifier.process_single_text(text, threshold),
-            inputs=[text_input, threshold_slider],
-            outputs=result_text
+                return f"Помилка: {str(e)}", classifier.get_cache_stats()
+        
+        def process_single_text(text, threshold):
+            """Обробка одного тексту"""
+            try:
+                return classifier.process_single_text(text, threshold)
+            except Exception as e:
+                return {"error": str(e)}
+        
+        def load_data(csv_path, emb_path):
+            """Завантаження даних для пакетної обробки"""
+            try:
+                return classifier.load_data(csv_path, emb_path)
+            except Exception as e:
+                return f"Помилка: {str(e)}"
+        
+        def classify_batch(filter_str, threshold):
+            """Пакетна класифікація"""
+            try:
+                return classifier.classify_rows(filter_str, threshold)
+            except Exception as e:
+                return None
+        
+        def save_results():
+            """Збереження результатів"""
+            try:
+                return classifier.save_results()
+            except Exception as e:
+                return f"Помилка: {str(e)}"
+        
+        # Підключення подій
+        model_type.change(
+            fn=update_model_inputs,
+            inputs=[model_type],
+            outputs=[model_choice, local_model_path, device_choice]
         )
         
         build_btn.click(
-            fn=update_with_file,
-            inputs=[json_file, model_choice, force_rebuild],
+            fn=update_classifier_settings,
+            inputs=[
+                json_file,
+                model_type,
+                model_choice,
+                local_model_path,
+                device_choice,
+                force_rebuild
+            ],
             outputs=[build_out, cache_stats]
         )
         
+        single_process_btn.click(
+            fn=process_single_text,
+            inputs=[text_input, threshold_slider],
+            outputs=result_text
+        )
+        
         load_btn.click(
-            fn=lambda csv, emb: classifier.load_data(csv, emb),
+            fn=load_data,
             inputs=[csv_input, emb_input],
             outputs=load_output
         )
         
         classify_btn.click(
-            fn=lambda filter_str, threshold: classifier.classify_rows(filter_str, threshold),
+            fn=classify_batch,
             inputs=[filter_in, batch_threshold],
             outputs=classify_out
         )
         
         save_btn.click(
-            fn=lambda: classifier.save_results("messages_with_labels.csv"),
+            fn=save_results,
             inputs=[],
             outputs=save_out
         )
 
+    # Запуск веб-інтерфейсу
     demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
 
 if __name__ == "__main__":

local_embedder.py

@@ -0,0 +1,113 @@
+import numpy as np
+import torch
+from typing import List, Union, Dict
+from transformers import AutoTokenizer, AutoModel
+from pathlib import Path
+import json
+
+class LocalEmbedder:
+    def __init__(self, model_name: str, device: str = None, batch_size: int = 32):
+        """
+        Ініціалізація локальної моделі для ембедінгів
+        
+        Args:
+            model_name: назва або шлях до моделі (з HuggingFace або локальна)
+            device: пристрій для обчислень ('cuda', 'cpu' або None - автовибір)
+            batch_size: розмір батчу для інференсу
+        """
+        self.model_name = model_name
+        self.batch_size = batch_size
+        
+        # Визначення пристрою
+        if device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = device
+            
+        # Завантаження моделі та токенізатора
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name).to(self.device)
+        self.model.eval()
+        
+        # Максимальна довжина послідовності
+        self.max_length = self.tokenizer.model_max_length
+        if self.max_length > 512:
+            self.max_length = 512
+            
+    def _normalize_embeddings(self, embeddings: np.ndarray) -> np.ndarray:
+        """
+        L2-нормалізація ембедінгів
+        
+        Args:
+            embeddings: матриця ембедінгів
+            
+        Returns:
+            np.ndarray: нормалізована матриця ембедінгів
+        """
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        return embeddings / norms
+            
+    def get_embeddings(self, texts: Union[str, List[str]]) -> np.ndarray:
+        """
+        Отримання ембедінгів для тексту або списку текстів
+        
+        Args:
+            texts: текст або список текстів
+            
+        Returns:
+            np.ndarray: матриця нормалізованих ембедінгів
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+            
+        all_embeddings = []
+        
+        with torch.no_grad():
+            for i in range(0, len(texts), self.batch_size):
+                batch_texts = texts[i:i + self.batch_size]
+                
+                # Токенізація
+                encoded = self.tokenizer.batch_encode_plus(
+                    batch_texts,
+                    padding=True,
+                    truncation=True,
+                    max_length=self.max_length,
+                    return_tensors='pt'
+                )
+                
+                # Переміщуємо тензори на потрібний пристрій
+                input_ids = encoded['input_ids'].to(self.device)
+                attention_mask = encoded['attention_mask'].to(self.device)
+                
+                # Отримуємо ембедінги
+                outputs = self.model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask
+                )
+                
+                # Використовуємо [CLS] токен як ембедінг
+                embeddings = outputs.last_hidden_state[:, 0, :]
+                all_embeddings.append(embeddings.cpu().numpy())
+                
+        # Об'єднуємо всі батчі
+        embeddings = np.vstack(all_embeddings)
+        
+        # Нормалізуємо ембедінги
+        normalized_embeddings = self._normalize_embeddings(embeddings)
+        
+        return normalized_embeddings
+    
+    def get_model_info(self) -> Dict[str, any]:
+        """
+        Отримання інформації про модель
+        
+        Returns:
+            Dict: інформація про модель
+        """
+        return {
+            'model_name': self.model_name,
+            'device': self.device,
+            'embedding_size': self.model.config.hidden_size,
+            'max_length': self.max_length,
+            'batch_size': self.batch_size
+        }

requirements.txt

@@ -3,4 +3,6 @@ openai
 pandas
 numpy
 python-dotenv
-scikit-learn
+scikit-learn
+torch
+transformers

sdc_classifier.py

@@ -2,21 +2,35 @@ import os
 import numpy as np
 import pandas as pd
 import json
-from typing import Dict, List
+from typing import Dict, List, Optional, Union
 from openai import OpenAI
 from pathlib import Path
 from embedding_cache import EmbeddingCache
 
 class SDCClassifier:
-    def __init__(self, openai_api_key: str = None, cache_path: str = "embeddings_cache.db"):
+    def __init__(self, 
+                 openai_api_key: str = None, 
+                 cache_path: str = "embeddings_cache.db",
+                 local_model: str = "cambridgeltl/SapBERT-from-PubMedBERT-fulltext",
+                 device: str = None):
         """
         Ініціалізація класифікатора SDC
         
         Args:
             openai_api_key: API ключ для OpenAI (опціонально, можна взяти з env)
             cache_path: шлях до файлу кешу ембедінгів
+            local_model: назва або шлях до локальної моделі
+            device: пристрій для локальної моделі ('cuda', 'cpu' або None)
         """
         self.client = OpenAI(api_key=openai_api_key or os.getenv("OPENAI_API_KEY"))
+        self.local_embedder = None
+        self.using_local = False
+        
+        if local_model:
+            from local_embedder import LocalEmbedder
+            self.local_embedder = LocalEmbedder(local_model, device)
+            self.using_local = True
+        
         self.classes_json = {}
         self.class_signatures = None
         self.df = None
@@ -24,13 +38,12 @@ class SDCClassifier:
         self.embeddings_mean = None
         self.embeddings_std = None
         
-        # Створення директорії для кешу, якщо потрібно
+        # Створення директорії для кешу
         cache_dir = os.path.dirname(cache_path)
         if cache_dir and not os.path.exists(cache_dir):
             os.makedirs(cache_dir)
             
         # Ініціалізація кешу
-        from embedding_cache import EmbeddingCache
         self.cache = EmbeddingCache(cache_path)
         
         # Базовий стан
@@ -66,60 +79,23 @@ class SDCClassifier:
         self.classes_json = self.base_classes_json.copy()
         self.class_signatures = self.base_signatures.copy() if self.base_signatures else None
 
-    def load_initial_state(self, classes_file: str, signatures_file: str) -> str:
+    def load_classes(self, json_path: Union[str, dict]) -> dict:
         """
-        Завантаження початкового стану при старті застосунку
+        Завантаження класів та їх хінтів з JSON файлу або словника
         
         Args:
-            classes_file: шлях до файлу з класами
-            signatures_file: шлях до файлу з signatures
+            json_path: шлях до JSON файлу або словник з класами
             
         Returns:
-            str: повідомлення про результат завантаження
+            dict: словник класів та їх хінтів
         """
         try:
-            self.base_classes_json = self.load_classes(classes_file)
-            if os.path.exists(signatures_file):
-                self.base_signatures = self.load_signatures(signatures_file)
-            
-            # Встановлюємо поточний стан як базовий
-            self.classes_json = self.base_classes_json.copy()
-            self.class_signatures = self.base_signatures.copy() if self.base_signatures else None
-            
-            return f"Завантажено {len(self.base_classes_json)} базових класів"
-        except Exception as e:
-            return f"Помилка при завантаженні базового стану: {str(e)}"
-            
-    def restore_base_state(self) -> None:
-        """Відновлення базового стану"""
-        self.classes_json = self.base_classes_json.copy()
-        self.class_signatures = self.base_signatures.copy() if self.base_signatures else None
-
-    def load_initial_state(self, classes_file: str, signatures_file: str):
-        """Завантаження початкового стану при старті застосунку"""
-        self.base_classes_json = self.load_classes(classes_file)
-        self.base_signatures = self.load_signatures(signatures_file)
-        
-        # Встановлюємо поточний стан як базовий
-        self.classes_json = self.base_classes_json.copy()
-        self.class_signatures = self.base_signatures.copy() if self.base_signatures else None
-        
-    def restore_base_state(self):
-        """Відновлення базового стану"""
-        self.classes_json = self.base_classes_json.copy()
-        self.class_signatures = self.base_signatures.copy() if self.base_signatures else None
-
-    def load_classes(self, json_path: str) -> dict:
-        """Завантаження класів та їх хінтів з JSON файлу"""
-        try:
-            # Якщо передано вміст файлу замість шляху
             if isinstance(json_path, dict):
                 self.classes_json = json_path
             else:
                 with open(json_path, 'r', encoding='utf-8') as f:
                     self.classes_json = json.load(f)
                     
-            # Валідація структури
             if not all(isinstance(hints, list) for hints in self.classes_json.values()):
                 raise ValueError("Кожен клас повинен мати список хінтів")
                 
@@ -132,12 +108,25 @@ class SDCClassifier:
             return {}
     
     def save_signatures(self, filename: str = "signatures.npz") -> None:
-        """Зберігає signatures у NPZ файл"""
+        """
+        Зберігає signatures у NPZ файл
+        
+        Args:
+            filename: шлях до файлу для збереження
+        """
         if self.class_signatures:
             np.savez(filename, **self.class_signatures)
 
     def load_signatures(self, filename: str = "signatures.npz") -> Dict[str, np.ndarray]:
-        """Завантажує signatures з NPZ файлу"""
+        """
+        Завантажує signatures з NPZ файлу
+        
+        Args:
+            filename: шлях до файлу з signatures
+            
+        Returns:
+            Dict[str, np.ndarray]: словник signatures
+        """
         try:
             with np.load(filename) as data:
                 self.class_signatures = {key: data[key] for key in data.files}
@@ -145,31 +134,34 @@ class SDCClassifier:
         except (FileNotFoundError, IOError):
             return None
 
-    def get_openai_embedding(self, text: str, model_name: str = "text-embedding-3-large") -> list:
+    def get_embedding(self, text: str, model_name: str = None) -> list:
         """
-        Отримання ембедінгу тексту через OpenAI API з використанням кешу
+        Отримання ембедінгу тексту
         
         Args:
             text: текст для ембедінгу
-            model_name: назва моделі OpenAI
+            model_name: назва моделі (OpenAI) або None для локальної
             
         Returns:
             list: ембедінг тексту
         """
-        # Спроба отримати з кешу
-        cached_embedding = self.cache.get(text, model_name)
+        # Перевіряємо кеш
+        cached_embedding = self.cache.get(text, model_name or "local")
         if cached_embedding is not None:
             return cached_embedding.tolist()
-            
-        # Якщо нема в кеші - отримуємо через API
-        response = self.client.embeddings.create(
-            input=text,
-            model=model_name
-        )
-        embedding = response.data[0].embedding
+        
+        # Отримуємо ембедінг
+        if self.using_local and model_name is None:
+            embedding = self.local_embedder.get_embeddings(text)[0]
+        else:
+            response = self.client.embeddings.create(
+                input=text,
+                model=model_name or "text-embedding-3-large"
+            )
+            embedding = response.data[0].embedding
         
         # Зберігаємо в кеш
-        self.cache.put(text, model_name, embedding)
+        self.cache.put(text, model_name or "local", embedding)
         
         return embedding
         
@@ -181,15 +173,24 @@ class SDCClassifier:
         """Очищення старих записів з кешу"""
         return self.cache.clear_old(days)
 
-    def embed_hints(self, hint_list: List[str], model_name: str) -> np.ndarray:
-        """Створення ембедінгів для списку хінтів"""
+    def embed_hints(self, hint_list: List[str], model_name: str = None) -> np.ndarray:
+        """
+        Створення ембедінгів для списку хінтів
+        
+        Args:
+            hint_list: список хінтів
+            model_name: назва моделі для ембедінгів
+            
+        Returns:
+            np.ndarray: матриця ембедінгів
+        """
         emb_list = []
         total_hints = len(hint_list)
         
         for idx, hint in enumerate(hint_list, 1):
             try:
                 print(f"  Отримання embedding {idx}/{total_hints}: '{hint}'")
-                emb = self.get_openai_embedding(hint, model_name=model_name)
+                emb = self.get_embedding(hint, model_name=model_name)
                 emb_list.append(emb)
             except Exception as e:
                 print(f"  Помилка при отриманні embedding для '{hint}': {str(e)}")
@@ -200,10 +201,10 @@ class SDCClassifier:
         
         return np.array(emb_list, dtype=np.float32)
 
-
-    def initialize_signatures(self, model_name: str = "text-embedding-3-large",
-                        signatures_file: str = "signatures.npz",
-                        force_rebuild: bool = False) -> str:
+    def initialize_signatures(self, 
+                            model_name: str = None,
+                            signatures_file: str = "signatures.npz",
+                            force_rebuild: bool = False) -> str:
         """
         Ініціалізує signatures: завантажує існуючі або створює нові
         
@@ -211,13 +212,16 @@ class SDCClassifier:
             model_name: назва моделі для ембедінгів
             signatures_file: шлях до файлу для збереження (None - не зберігати)
             force_rebuild: примусово перебудувати signatures
+            
+        Returns:
+            str: повідомлення про результат
         """
         if not self.classes_json:
             return "Помилка: Не знайдено жодного класу в classes.json"
         
         print(f"Знайдено {len(self.classes_json)} класів")
         
-        # Завантажуємо існуючі signatures, якщо є файл і не примусове оновлення
+        # Завантажуємо існуючі signatures
         if not force_rebuild and signatures_file and os.path.exists(signatures_file):
             try:
                 loaded_signatures = self.load_signatures(signatures_file)
@@ -250,7 +254,7 @@ class SDCClassifier:
             if not self.class_signatures:
                 return "Помилка: Не вдалося створити жодного signature"
             
-            # Зберігаємо signatures тільки якщо вказано шлях до файлу
+            # Зберігаємо signatures
             if signatures_file:
                 try:
                     self.save_signatures(signatures_file)
@@ -262,8 +266,17 @@ class SDCClassifier:
         except Exception as e:
             return f"Помилка при створенні signatures: {str(e)}"
     
-    def load_data(self, csv_path: str = "messages.csv", emb_path: str = "embeddings.npy"):
-        """Завантаження даних з CSV та NPY файлів"""
+    def load_data(self, csv_path: str = "messages.csv", emb_path: str = "embeddings.npy") -> str:
+        """
+        Завантаження даних з CSV та NPY файлів
+        
+        Args:
+            csv_path: шлях до CSV файлу
+            emb_path: шлях до NPY файлу з ембедінгами
+            
+        Returns:
+            str: повідомлення про результат
+        """
         self.df = pd.read_csv(csv_path)
         emb_local = np.load(emb_path)
         assert len(self.df) == len(emb_local), "CSV і embeddings різної довжини!"
@@ -277,7 +290,16 @@ class SDCClassifier:
         return f"Завантажено {len(self.df)} рядків"
 
     def predict_classes(self, text_embedding: np.ndarray, threshold: float = 0.0) -> Dict[str, float]:
-        """Передбачення класів для одного тексту"""
+        """
+        Передбачення класів для одного тексту
+        
+        Args:
+            text_embedding: ембедінг тексту
+            threshold: поріг впевненості
+            
+        Returns:
+            Dict[str, float]: словник класів та їх scores
+        """
         results = {}
         for cls, sign in self.class_signatures.items():
             score = float(np.dot(text_embedding, sign))
@@ -287,11 +309,20 @@ class SDCClassifier:
         return dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
 
     def process_single_text(self, text: str, threshold: float = 0.3) -> dict:
-        """Обробка одного тексту"""
+        """
+        Обробка одного тексту
+        
+        Args:
+            text: текст для класифікації
+            threshold: поріг впевненості
+            
+        Returns:
+            dict: результати класифікації
+        """
         if self.class_signatures is None:
             return {"error": "Спочатку збудуйте signatures!"}
             
-        emb = self.get_openai_embedding(text)
+        emb = self.get_embedding(text)
         
         if self.embeddings_mean is not None and self.embeddings_std is not None:
             emb = (emb - self.embeddings_mean) / self.embeddings_std
@@ -310,17 +341,26 @@ class SDCClassifier:
             "result": "\n".join(formatted_results)
         }
 
-    def classify_rows(self, filter_substring: str = "", threshold: float = 0.3):
-        """Класифікація всіх або відфільтрованих рядків"""
+    def classify_rows(self, filter_substring: str = "", threshold: float = 0.3) -> pd.DataFrame:
+        """
+        Класифікація всіх або відфільтрованих рядків
+        
+        Args:
+            filter_substring: підрядок для фільтрації
+            threshold: поріг впевненості
+            
+        Returns:
+            pd.DataFrame: результати класифікації
+        """
         if self.class_signatures is None:
-            return "Спочатку збудуйте signatures!"
+            raise ValueError("Спочатку збудуйте signatures!")
 
         if self.df is None or self.embeddings is None:
-            return "Дані не завантажені! Спочатку викличте load_data."
+            raise ValueError("Дані не завантажені! Спочатку викличте load_data.")
 
         if filter_substring:
             filtered_idx = self.df[self.df["Message"].str.contains(filter_substring, 
-                                                                 case=False, 
+                                                                 case=False,
                                                                  na=False)].index
         else:
             filtered_idx = self.df.index
@@ -345,9 +385,51 @@ class SDCClassifier:
         return result_df.reset_index(drop=True)
 
     def save_results(self, output_path: str = "messages_with_labels.csv") -> str:
-        """Зберігання результатів класифікації"""
+        """
+        Зберігання результатів класифікації
+        
+        Args:
+            output_path: шлях для збереження результатів
+            
+        Returns:
+            str: повідомлення про результат
+        """
         if self.df is None:
             return "Дані відсутні!"
         
         self.df.to_csv(output_path, index=False)
-        return f"Дані збережено у файл {output_path}"
+        return f"Дані збережено у файл {output_path}"
+        
+    def save_model_info(self, path: str = "model_info.json") -> None:
+        """
+        Зберігання інформації про поточний стан моделі
+        
+        Args:
+            path: шлях для збереження
+        """
+        info = {
+            "using_local": self.using_local,
+            "classes_count": len(self.classes_json),
+            "signatures_count": len(self.class_signatures) if self.class_signatures else 0,
+            "cache_stats": self.get_cache_stats(),
+        }
+        
+        if self.using_local:
+            info["local_model"] = self.local_embedder.get_model_info()
+            
+        with open(path, 'w', encoding='utf-8') as f:
+            json.dump(info, f, indent=2)
+            
+    @staticmethod
+    def load_model_info(path: str) -> dict:
+        """
+        Завантаження інформації про модель
+        
+        Args:
+            path: шлях до файлу з інформацією
+            
+        Returns:
+            dict: інформація про модель
+        """
+        with open(path, 'r', encoding='utf-8') as f:
+            return json.load(f)