23ws-LLMcoder/LLMcoder-GitHub-Python-Mix-Direct

1import os2import shutil3import time4from typing import Tuple5 6import numpy as np7import pandas as pd8import pytest9from sklearn.pipeline import Pipeline10 11from pmtrendviz.models.factory import ModelFactory12from pmtrendviz.models.manager import SimpleSklearnPipelineManager13from pmtrendviz.train.data import sample_training_data14from pmtrendviz.utils import get_es_client, get_models_dir15 16 17@pytest.mark.usefixtures('manage_pytest_index', 'manage_prediction_indices')18class TestSimpleSklearnPipelineManager:19    sample_data: pd.DataFrame = pd.DataFrame()20    n_samples: int = 100021    index: str = 'pytest'22    model_name: str = 'pytest-model'23    model = 'tfidf_truncatedsvd_kmeans'24    stop_words = 'english'25    max_df: float = 0.99926    min_df: float = 0.00127    n_components: int = 10028    n_clusters: int = 6929    ngram_range: Tuple[int, int] = (1, 1)30    random_state: int = 4231    include_title: bool = True32    include_abstract: bool = True33    include_keywords_major: bool = False34    include_keywords_minor: bool = False35 36    @classmethod37    def setup_class(cls) -> None:38        # Clear the test model directory39        shutil.rmtree(os.path.join(get_models_dir(), cls.model_name), ignore_errors=True)40 41        # Sample some data42        cls.sample_data = sample_training_data(cls.index, random_state=42, n_samples=cls.n_samples, method='forward')['text']43 44    @classmethod45    def teardown_class(cls) -> None:46        # Clear the test model directory47        shutil.rmtree(os.path.join(get_models_dir(), cls.model_name), ignore_errors=True)48 49    def teardown_method(self) -> None:50        # Clear the test model directory51        shutil.rmtree(os.path.join(get_models_dir(), self.model_name), ignore_errors=True)52 53    def test_sample_data(self) -> None:54        assert len(self.sample_data) == self.n_samples55 56    def test_create_default_simple_sklearn_pipeline_manager(self) -> None:57        manager = ModelFactory.create(58            self.model,59            self.include_title,60            self.include_abstract,61            self.include_keywords_major,62            self.include_keywords_minor,63            stop_words=self.stop_words,64            max_df=self.max_df,65            min_df=self.min_df,66            n_components=self.n_components,67            n_clusters=self.n_clusters,68            ngram_range=self.ngram_range,69            random_state=None)70 71        assert manager is not None72        assert isinstance(manager.model, Pipeline)73        assert manager.random_state is None74        assert manager.preprocessing_args == {75            'include_title': True,76            'include_abstract': True,77            'include_keywords_major': False,78            'include_keywords_minor': False,79        }80 81        assert (manager.cluster_names.columns == ['name']).all()82        assert len(manager.cluster_names) == 083 84    def test_create_custom_simple_sklearn_pipeline_manager(self) -> None:85        manager = ModelFactory.create(86            self.model,87            self.include_title,88            self.include_abstract,89            True,90            self.include_keywords_minor,91            stop_words=self.stop_words,92            max_df=self.max_df,93            min_df=self.min_df,94            n_components=self.n_components,95            n_clusters=self.n_clusters,96            ngram_range=self.ngram_range,97            random_state=42)98 99        assert manager.random_state == 42100        assert manager.preprocessing_args == {101            'include_title': True,102            'include_abstract': True,103            'include_keywords_major': True,104            'include_keywords_minor': False,105        }106 107    def test_set_name(self) -> None:108        manager = ModelFactory.create(109            self.model,110            self.include_title,111            self.include_abstract,112            self.include_keywords_major,113            self.include_keywords_minor,114            stop_words=self.stop_words,115            max_df=self.max_df,116            min_df=self.min_df,117            n_components=self.n_components,118            n_clusters=self.n_clusters,119            ngram_range=self.ngram_range,120            random_state=self.random_state)121        manager.set_name(self.model_name)122 123        assert manager._name == self.model_name124 125    def test_construct_predictions_index_name(self) -> None:126        manager = ModelFactory.create(127            self.model,128            self.include_title,129            self.include_abstract,130            self.include_keywords_major,131            self.include_keywords_minor,132            stop_words=self.stop_words,133            max_df=self.max_df,134            min_df=self.min_df,135            n_components=self.n_components,136            n_clusters=self.n_clusters,137            ngram_range=self.ngram_range,138            random_state=self.random_state)139        manager.set_name(self.model_name)140        predictions_index_name = manager._construct_predictions_index_name(self.index)141        predictions_file_name = manager._construct_predictions_file_name(self.index)142 143        assert manager._name == self.model_name144        assert predictions_index_name == f'{self.index}_{manager._name}_predictions'145        assert predictions_file_name == f'{self.index}_predictions.json'146 147    def test_fit(self) -> None:148        manager = ModelFactory.create(149            self.model,150            self.include_title,151            self.include_abstract,152            self.include_keywords_major,153            self.include_keywords_minor,154            stop_words=self.stop_words,155            max_df=self.max_df,156            min_df=self.min_df,157            n_components=self.n_components,158            n_clusters=self.n_clusters,159            ngram_range=self.ngram_range,160            random_state=self.random_state)161 162        manager.fit(self.sample_data)163 164        assert manager.model['clustering_kmeans'].labels_.shape == (self.n_samples,)165        assert len(manager.cluster_names) == self.n_clusters166 167    def test_predict(self) -> None:168        manager = ModelFactory.create(169            self.model,170            self.include_title,171            self.include_abstract,172            self.include_keywords_major,173            self.include_keywords_minor,174            stop_words=self.stop_words,175            max_df=self.max_df,176            min_df=self.min_df,177            n_components=self.n_components,178            n_clusters=self.n_clusters,179            ngram_range=self.ngram_range,180            random_state=self.random_state)181 182        manager.fit(self.sample_data)183        predictions = manager.predict(self.sample_data)184 185        assert isinstance(predictions, np.ndarray)186        assert