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

1le_data: pd.DataFrame = pd.DataFrame()2    n_samples: int = 10003    index: str = 'pytest'4    model_name: str = 'pytest-model'5    model = 'tfidf_truncatedsvd_kmeans'6    stop_words = 'english'7    max_df: float = 0.9998    min_df: float = 0.0019    n_components: int = 10010    n_clusters: int = 6911    ngram_range: Tuple[int, int] = (1, 1)12    random_state: int = 4213    include_title: bool = True14    include_abstract: bool = True15    include_keywords_major: bool = False16    include_keywords_minor: bool = False17 18    @classmethod19    def setup_class(cls) -> None:20        # Clear the test model directory21        shutil.rmtree(os.path.join(get_models_dir(), cls.model_name), ignore_errors=True)22 23        # Sample some data24        cls.sample_data = sample_training_data(cls.index, random_state=42, n_samples=cls.n_samples, method='forward')['text']25 26    @classmethod27    def teardown_class(cls) -> None:28        # Clear the test model directory29        shutil.rmtree(os.path.join(get_models_dir(), cls.model_name), ignore_errors=True)30 31    def teardown_method(self) -> None:32        # Clear the test model directory33        shutil.rmtree(os.path.join(get_models_dir(), self.model_name), ignore_errors=True)34 35    def test_sample_data(self) -> None:36        assert len(self.sample_data) == self.n_samples37 38    def test_create_default_simple_sklearn_pipeline_manager(self) -> None:39        manager = ModelFactory.create(40            self.model,41            self.include_title,42            self.include_abstract,43            self.include_keywords_major,44            self.include_keywords_minor,45            stop_words=self.stop_words,46            max_df=self.max_df,47            min_df=self.min_df,48            n_components=self.n_components,49            n_clusters=self.n_clusters,50            ngram_range=self.ngram_range,51            random_state=None)52 53        assert manager is not None54        assert isinstance(manager.model, Pipeline)55        assert manager.random_state is None56        assert manager.preprocessing_args == {57            'include_title': True,58            'include_abstract': True,59            'include_keywords_major': False,60            'include_keywords_minor': False,61        }62 63        assert (manager.cluster_names.columns == ['name']).all()64        assert len(manager.cluster_names) == 065 66    def test_create_custom_simple_sklearn_pipeline_manager(self) -> None:67        manager = ModelFactory.create(68            self.model,69            self.include_title,70            self.include_abstract,71            True,72            self.include_keywords_minor,73            stop_words=self.stop_words,74            max_df=self.max_df,75            min_df=self.min_df,76            n_components=self.n_components,77            n_clusters=self.n_clusters,78            ngram_range=self.ngram_range,79            random_state=42)80 81        assert manager.random_state == 4282        assert manager.preprocessing_args == {83            'include_title': True,84            'include_abstract': True,85            'include_keywords_major': True,86            'include_keywords_minor': False,87        }88 89    def test_set_name(self) -> None:90        manager = ModelFactory.create(91            self.model,92            self.include_title,93            self.include_abstract,94            self.include_keywords_major,95            self.include_keywords_minor,96            stop_words=self.stop_words,97            max_df=self.max_df,98            min_df=self.min_df,99            n_components=self.n_components,100            n_clusters=self.n_clusters,101            ngram_range=self.ngram_range,102            random_state=self.random_state)103        manager.set_name(self.model_name)104 105        assert manager._name == self.model_name106 107    def test_construct_predictions_index_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        predictions_index_name = manager._construct_predictions_index_name(self.index)123        predictions_file_name = manager._construct_predictions_file_name(self.index)124 125        assert manager._name == self.model_name126        assert predictions_index_name == f'{self.index}_{manager._name}_predictions'127        assert predictions_file_name == f'{self.index}_predictions.json'128 129    def test_fit(self) -> None:130        manager = ModelFactory.create(131            self.model,132            self.include_title,133            self.include_abstract,134            self.include_keywords_major,135            self.include_keywords_minor,136            stop_words=self.stop_words,137            max_df=self.max_df,138            min_df=self.min_df,139            n_components=self.n_components,140            n_clusters=self.n_clusters,141            ngram_range=self.ngram_range,142            random_state=self.random_state)143 144        manager.fit(self.sample_data)145 146        assert manager.model['clustering_kmeans'].labels_.shape == (self.n_samples,)147        assert len(manager.cluster_names) == self.n_clusters148 149    def test_predict(self) -> None:150        manager = ModelFactory.create(151            self.model,152            self.include_title,153            self.include_abstract,154            self.include_keywords_major,155            self.include_keywords_minor,156            stop_words=self.stop_words,157            max_df=self.max_df,158            min_df=self.min_df,159            n_components=self.n_components,160            n_clusters=self.n_clusters,161            ngram_range=self.ngram_range,162            random_state=self.random_state)163 164        manager.fit(self.sample_data)165        predictions = manager.predict(self.sample_data)166 167        assert isinstance(predictions, np.ndarray)168        assert predictions.shape == (self.n_samples,)169 170    def test_transform(self) -> None:171        manager = ModelFactory.create(172            self.model,173            self.include_title,174            self.include_abstract,175            self.include_keywords_major,176            self.include_keywords_minor,177            stop_words=self.stop_words,178            max_df=self.max_df,179            min_df=self.min_df,180            n_components=self.n_components,181            n_clusters=self.n_clusters,182            ngram_range=self.ngram_range,183            random_state=self.random_state)184 185        manager.fit(self.sample_data)186        transformed = manager.transform(self.sample_data, exclude_from='clustering')187 188        assert isinstance(transformed, np.ndarray)189        assert transformed.shape == (self.n_samples, self.n_components)190 191        transformed_with_kmeans = manager.transform(self.sample_data)192 193        assert isinstance(transformed_with_kmeans, np.ndarray)194        assert transformed_with_kmeans.shape == (self.n_samples, self.n_clusters)195 196    def test_save(self) -> None:197        manager = ModelFactory.create(198            self.model,199            self.include_title,200            self.include_abstract,201            self.include_keywords_major,202            self.include_keywords_minor,203            stop_words=self.stop_words,204            max_df=self.max_df,205            min_df=self.min_df,206            n_components=self.n_components,207            n_clusters=self.n_clusters,208            ngram_range=self.ngram_range,209            random_state=self.random_state)210 211        manager.fit(self.sample_data)212        manager.save(self.model_name)213 214        assert os.path.exists(os.path.join(get_models_dir(), self.model_name, 'model.joblib'))215        assert os.path.exists(os.path.join(get_models_dir(), self.model_name, 'config.json'))216        assert os.path.exists(os.path.join(get_models_dir(), self.model_name, 'cluster_names.csv'))217 218    def test_save_only_model(self) -> None:219        manager = ModelFactory.create(220            self.model,221            self.include_title,222            self.include_abstract,223            self.include_keywords_major,224            self.include_keywords_minor,225            stop_words=self.stop_words,226            max_df=self.max_df,227            min_df=self.min_df,228            n_components=self.n_components,229            n_clusters=self.n_clusters,230            ngram_range=self.ngram_range,231            random_state=self.random_state)232 233        manager.fit(self.sample_data)234        manager.save(self.model_name, save_only='model')235 236        assert os.path.exists(os.path.join(get_models_dir(), self.model_name, 'model.joblib'))237        assert not os.path.exists(os.path.join(get_models_dir(), self.model_name, 'config.json'))238        assert not os.path.exists(os.path.join(get_models_dir(), self.model_name, 'cluster_names.csv'))239 240    def test_load(self) -> None:241        manager = ModelFactory.create(242            self.model,243            self.include_title,244            self.include_abstract,245            True,246            self.include_keywords_minor,247            stop_words=self.stop_words,248            max_df=self.max_df,249            min_df=self.min_df,250            n_components=self.n_components,251            n_clusters=self.n_clusters,252            ngram_range=self.ngram_range,253            random_state=self.random_state)254 255        manager.fit(self.sample_data)256        manager.save(self.model_name)257 258        manager_predictions = manager.model.predict(self.sample_data)259        manager_transformed = manager.model.transform(self.sample_data)260 261        loaded_manager = SimpleSklearnPipelineManager.load(self