classifier_app.py

from sdc_classifier import SDCClassifier
from dotenv import load_dotenv
import torch
import json
import os
from typing import Dict, Tuple, Optional, Any, List
from dataclasses import dataclass, field

# Load environment variables
load_dotenv()

@dataclass
class Config:
    DEFAULT_CLASSES_FILE: str = "classes.json"
    DEFAULT_SIGNATURES_FILE: str = "signatures.npz"
    CACHE_FILE: str = "embeddings_cache.db"
    MODEL_INFO_FILE: str = "model_info.json"
    DEFAULT_OPENAI_MODELS: List[str] = field(default_factory=lambda: ["text-embedding-3-large", "text-embedding-3-small"])
    DEFAULT_LOCAL_MODEL: str = "cambridgeltl/SapBERT-from-PubMedBERT-fulltext"

config = Config()

class ClassifierApp:
    def __init__(self):
        self.classifier = None
        self.initial_info = {
            "status": "initializing",
            "model_info": {},
            "classes_info": {},
            "errors": []
        }

    def initialize_environment(self) -> Tuple[Dict, Optional[SDCClassifier]]:
        """Ініціалізація середовища при першому запуску"""
        try:
            # Перевіряємо наявність необхідних файлів
            if not os.path.exists(config.DEFAULT_CLASSES_FILE):
                self.initial_info["errors"].append(
                    f"ПОМИЛКА: Файл {config.DEFAULT_CLASSES_FILE} не знайдено!"
                )
                self.initial_info["status"] = "error"
                print(f"\nПомилка: Файл {config.DEFAULT_CLASSES_FILE} не знайдено!")
                return self.initial_info, None
                
            print("\nСтворення класифікатора...")
            try:
                # Визначаємо яка модель використовувалась для сигнатур
                signatures_model = None
                if os.path.exists(config.MODEL_INFO_FILE):
                    with open(config.MODEL_INFO_FILE, 'r') as f:
                        model_info = json.load(f)
                        if not model_info.get('using_local', True):
                            signatures_model = "text-embedding-3-small"  # Модель, яка використовувалась
                
                # Створюємо класифікатор з тією ж моделлю
                self.classifier = SDCClassifier(openai_api_key=os.getenv("OPENAI_API_KEY"))
                print(f"Використовується модель: {signatures_model or 'local'}")
                
            except Exception as e:
                print(f"\nПомилка при створенні класифікатора: {str(e)}")
                self.initial_info["errors"].append(f"Помилка при створенні класифікатора: {str(e)}")
                self.initial_info["status"] = "error"
                return self.initial_info, None
            
            print("\nЗавантаження класів...")
            try:
                classes = self.classifier.load_classes(config.DEFAULT_CLASSES_FILE)
                self.initial_info["classes_info"] = {
                    "total_classes": len(classes),
                    "classes_list": list(classes.keys()),
                    "hints_per_class": {cls: len(hints) for cls, hints in classes.items()}
                }
            except Exception as e:
                print(f"\nПомилка при завантаженні класів: {str(e)}")
                self.initial_info["errors"].append(f"Помилка при завантаженні класів: {str(e)}")
                self.initial_info["status"] = "error"
                return self.initial_info, None
            
            print("\nПеревірка та завантаження сигнатур...")
            if os.path.exists(config.DEFAULT_SIGNATURES_FILE):
                try:
                    self.classifier.load_signatures(config.DEFAULT_SIGNATURES_FILE)
                    self.initial_info["status"] = "success"
                    print("Сигнатури завантажено успішно")
                except Exception as e:
                    print(f"\nПомилка при завантаженні сигнатур: {str(e)}")
                    self.initial_info["errors"].append(f"Помилка при завантаженні сигнатур: {str(e)}")
                    self.initial_info["status"] = "error"
                    return self.initial_info, None
            else:
                print("\nСтворення нових сигнатур...")
                self.initial_info["status"] = "creating_signatures"
                try:
                    result = self.classifier.initialize_signatures(
                        force_rebuild=True,
                        signatures_file=config.DEFAULT_SIGNATURES_FILE
                    )
                    if isinstance(result, str) and "error" in result.lower():
                        print(f"\nПомилка при створенні сигнатур: {result}")
                        self.initial_info["errors"].append(result)
                        self.initial_info["status"] = "error"
                        return self.initial_info, None
                except Exception as e:
                    print(f"\nПомилка при створенні сигнатур: {str(e)}")
                    self.initial_info["errors"].append(f"Помилка при створенні сигнатур: {str(e)}")
                    self.initial_info["status"] = "error"
                    return self.initial_info, None
            
            print("\nЗбереження інформації про модель...")
            try:
                self.classifier.save_model_info(config.MODEL_INFO_FILE)
                with open(config.MODEL_INFO_FILE, "r") as f:
                    self.initial_info["model_info"] = json.load(f)
                    
                self.initial_info["status"] = "success"
                print("\nІніціалізація завершена успішно")
                return self.initial_info, self.classifier
                    
            except Exception as e:
                print(f"\nПомилка при збереженні інформації про модель: {str(e)}")
                self.initial_info["errors"].append(f"Помилка при читанні model_info: {str(e)}")
                self.initial_info["status"] = "error"
                return self.initial_info, None
                    
        except Exception as e:
            print(f"\nЗагальна помилка при ініціалізації: {str(e)}")
            self.initial_info["errors"].append(f"ПОМИЛКА при ініціалізації: {str(e)}")
            self.initial_info["status"] = "error"
            return self.initial_info, None

    def create_classifier(
        self,
        model_type: str,
        openai_model: Optional[str] = None,
        local_model: Optional[str] = None,
        device: Optional[str] = None
    ) -> SDCClassifier:
        """Створення класифікатора з відповідними параметрами"""
        classifier = SDCClassifier()
        
        if model_type == "OpenAI":
            if hasattr(classifier, 'set_openai_model'):
                classifier.set_openai_model(openai_model)
        else:
            if hasattr(classifier, 'set_local_model'):
                classifier.set_local_model(local_model, device)
        
        return classifier

    def update_model_inputs(
        self,
        model_type: str,
        openai_model: str,
        local_model: str,
        device: str
    ) -> Dict[str, Any]:
        """Оновлення моделі та інтерфейсу при зміні типу моделі"""
        try:
            self.classifier = self.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
            )
            
            self.classifier.restore_base_state()
            result = self.classifier.initialize_signatures()
            
            self.classifier.save_model_info(config.MODEL_INFO_FILE)
            with open(config.MODEL_INFO_FILE, "r") as f:
                model_info = json.load(f)
            
            new_system_info = {
                "status": "success",
                "model_info": model_info,
                "classes_info": {
                    "total_classes": len(self.classifier.classes_json),
                    "classes_list": list(self.classifier.classes_json.keys()),
                    "hints_per_class": {cls: len(hints) for cls, hints in self.classifier.classes_json.items()}
                },
                "errors": []
            }
            
            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"),
                "system_info": new_system_info,
                "system_md": self.update_system_markdown(new_system_info),
                "build_out": f"Модель змінено на {model_type}",
                "cache_stats": self.classifier.get_cache_stats()
            }
            
        except Exception as e:
            error_info = {
                "status": "error",
                "errors": [str(e)],
                "model_info": {},
                "classes_info": {}
            }
            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"),
                "system_info": error_info,
                "system_md": self.update_system_markdown(error_info),
                "build_out": f"Помилка: {str(e)}",
                "cache_stats": {}
            }

    def update_classifier_settings(
        self,
        json_file: Optional[str],
        model_type: str,
        openai_model: str,
        local_model: str,
        device: str,
        force_rebuild: bool
    ) -> Tuple[str, Dict, Dict, str]:
        """Оновлення налаштувань класифікатора"""
        try:
            self.classifier = self.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
            )
            
            if json_file is not None:
                with open(json_file.name, 'r', encoding='utf-8') as f:
                    new_classes = json.load(f)
                self.classifier.load_classes(new_classes)
            else:
                self.classifier.restore_base_state()
            
            result = self.classifier.initialize_signatures(
                force_rebuild=force_rebuild,
                signatures_file=config.DEFAULT_SIGNATURES_FILE if not force_rebuild else None
            )
            
            self.classifier.save_model_info(config.MODEL_INFO_FILE)
            with open(config.MODEL_INFO_FILE, "r") as f:
                model_info = json.load(f)
            
            new_system_info = {
                "status": "success",
                "model_info": model_info,
                "classes_info": {
                    "total_classes": len(self.classifier.classes_json),
                    "classes_list": list(self.classifier.classes_json.keys()),
                    "hints_per_class": {
                        cls: len(hints) 
                        for cls, hints in self.classifier.classes_json.items()
                    }
                },
                "errors": []
            }
            
            return (
                result,
                self.classifier.get_cache_stats(),
                new_system_info,
                self.update_system_markdown(new_system_info)
            )
        except Exception as e:
            error_info = {
                "status": "error",
                "errors": [str(e)],
                "model_info": {},
                "classes_info": {}
            }
            return (
                f"Помилка: {str(e)}",
                self.classifier.get_cache_stats(),
                error_info,
                self.update_system_markdown(error_info)
            )

    def process_single_text(self, text: str, threshold: float) -> Dict:
        """Обробка одного тексту"""
        try:
            if self.classifier is None:
                raise ValueError("Класифікатор не ініціалізовано")
            return self.classifier.process_single_text(text, threshold)
        except Exception as e:
            return {"error": str(e)}

    def load_data(self, csv_path: str, emb_path: str) -> str:
        """Завантаження даних для пакетної обробки"""
        try:
            if self.classifier is None:
                raise ValueError("Класифікатор не ініціалізовано")
            return self.classifier.load_data(csv_path, emb_path)
        except Exception as e:
            return f"Помилка: {str(e)}"

    def classify_batch(self, filter_str: str, threshold: float):
        """Пакетна класифікація"""
        try:
            if self.classifier is None:
                raise ValueError("Класифікатор не ініціалізовано")
            return self.classifier.classify_rows(filter_str, threshold)
        except Exception as e:
            return None

    def save_results(self) -> str:
        """Збереження результатів"""
        try:
            if self.classifier is None:
                raise ValueError("Класифікатор не ініціалізовано")
            return self.classifier.save_results()
        except Exception as e:
            return f"Помилка: {str(e)}"

    def sync_system_info(self) -> Dict:
        """Синхронізація системної інформації"""
        try:
            if self.classifier is None:
                raise ValueError("Класифікатор не ініціалізовано")
                
            self.classifier.save_model_info(config.MODEL_INFO_FILE)
            with open(config.MODEL_INFO_FILE, "r") as f:
                model_info = json.load(f)
            
            self.initial_info = {
                "status": "success",
                "model_info": model_info,
                "classes_info": {
                    "total_classes": len(self.classifier.classes_json),
                    "classes_list": list(self.classifier.classes_json.keys()),
                    "hints_per_class": {
                        cls: len(hints) 
                        for cls, hints in self.classifier.classes_json.items()
                    }
                },
                "errors": []
            }
            return self.initial_info
        except Exception as e:
            self.initial_info = {
                "status": "error",
                "model_info": {},
                "classes_info": {},
                "errors": [str(e)]
            }
            return self.initial_info

    @staticmethod
    def update_system_markdown(info: Dict) -> str:
        """Оновлення Markdown з системною інформацією"""
        if info["status"] == "success":
            return f"""
            ### Поточна конфігурація:
            - Модель: {info['model_info'].get('using_local', 'OpenAI')}
            - Кількість класів: {info['classes_info']['total_classes']}
            - Класи: {', '.join(info['classes_info']['classes_list'])}
            """
        else:
            return f"""
            ### Помилки ініціалізації:
            {chr(10).join('- ' + err for err in info['errors'])}
            """