bertrend 0.3.1__tar.gz → 0.3.3__tar.gz

{bertrend-0.3.1 → bertrend-0.3.3}/PKG-INFO

@@ -1,90 +1,88 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: bertrend
-Version: 0.3.1
+Version: 0.3.3
 Summary: BERTrend analyses topic evolution over time using state-of-the-art transformer models: topic modelling and LLM
-Home-page: https://github.com/rte-france/BERTrend
 License: LICENSE.md
-Keywords: topics,BERTopic,transformers,weak signals
+Keywords: topics,BERTopic,transformers,weak signals,LLM
 Author: Jérôme Picault
 Author-email: jerome.picault@rte-france.com
+Maintainer: Jérôme Picault
+Maintainer-email: jerome.picault@rte-france.com
 Requires-Python: >=3.12,<4.0
-Classifier: License :: Other/Proprietary License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Provides-Extra: apps
-Provides-Extra: tests
 Requires-Dist: Pyqtree (==1.0.0)
-Requires-Dist: arxiv (>=2.2.0,<3.0.0) ; extra == "apps"
+Requires-Dist: arxiv (>=2.2.0) ; extra == "apps"
 Requires-Dist: bertopic (==0.16.2)
-Requires-Dist: black (>=24.10.0,<25.0.0)
-Requires-Dist: coverage (>=7.6.9,<8.0.0) ; extra == "tests"
-Requires-Dist: cron-descriptor (>=1.4.5,<2.0.0)
+Requires-Dist: black (>=24.10.0)
+Requires-Dist: coverage (>=7.6.9)
+Requires-Dist: cron-descriptor (>=1.4.5)
 Requires-Dist: dask (==2024.12.0)
 Requires-Dist: datamapplot (==0.3.0)
-Requires-Dist: dateparser (>=1.2.0,<2.0.0)
-Requires-Dist: dill (>=0.4.0,<0.5.0)
-Requires-Dist: feedparser (>=6.0.11,<7.0.0) ; extra == "apps"
-Requires-Dist: gensim (==4.3.2)
-Requires-Dist: google-api-python-client (>=2.148.0,<3.0.0) ; extra == "apps"
-Requires-Dist: google-auth-oauthlib (>=1.2.1,<2.0.0) ; extra == "apps"
-Requires-Dist: googlenewsdecoder (>=0.1.7,<0.2.0) ; extra == "apps"
-Requires-Dist: goose3 (>=3.1.19,<4.0.0) ; extra == "apps"
-Requires-Dist: hdbscan (>=0.8.40,<0.9.0)
-Requires-Dist: ipython (>=8.28.0,<9.0.0)
-Requires-Dist: joblib (>=1.4.2,<2.0.0)
-Requires-Dist: jsonlines (>=4.0.0,<5.0.0)
-Requires-Dist: langchain-chroma (>=0.2.2,<0.3.0)
-Requires-Dist: langchain-community (>=0.3.21,<0.4.0)
-Requires-Dist: langchain-openai (>=0.3.12,<0.4.0)
-Requires-Dist: langdetect (>=1.0.9,<2.0.0)
-Requires-Dist: loguru (>=0.7.3,<0.8.0)
-Requires-Dist: lxml-html-clean (>=0.2.2,<0.3.0)
-Requires-Dist: markdown (>=3.7,<4.0)
-Requires-Dist: newscatcherapi (>=0.7.3,<0.8.0) ; extra == "apps"
-Requires-Dist: newspaper4k (>=0.9.3.1,<0.10.0.0) ; extra == "apps"
-Requires-Dist: nltk (>=3.9.1,<4.0.0)
+Requires-Dist: dateparser (>=1.2.0)
+Requires-Dist: dill (>=0.4.0)
+Requires-Dist: feedparser (>=6.0.11) ; extra == "apps"
+Requires-Dist: gensim (>=4.3.3)
+Requires-Dist: google-api-python-client (>=2.148.0) ; extra == "apps"
+Requires-Dist: google-auth-oauthlib (>=1.2.1) ; extra == "apps"
+Requires-Dist: googlenewsdecoder (>=0.1.7) ; extra == "apps"
+Requires-Dist: goose3 (>=3.1.19) ; extra == "apps"
+Requires-Dist: hdbscan (>=0.8.40)
+Requires-Dist: ipython (>=8.28.0)
+Requires-Dist: joblib (>=1.4.2)
+Requires-Dist: jsonlines (>=4.0.0)
+Requires-Dist: langchain-core (>=0.3.59)
+Requires-Dist: langdetect (>=1.0.9)
+Requires-Dist: loguru (>=0.7.3)
+Requires-Dist: lxml-html-clean (>=0.2.2)
+Requires-Dist: markdown (>=3.7)
+Requires-Dist: newscatcherapi (>=0.7.3) ; extra == "apps"
+Requires-Dist: newspaper4k (>=0.9.3.1) ; extra == "apps"
+Requires-Dist: nltk (>=3.9.1)
 Requires-Dist: numpy (<2)
-Requires-Dist: openai (>=1.77.0,<2.0.0)
+Requires-Dist: openai (>=1.77.0)
 Requires-Dist: opentelemetry-exporter-otlp-proto-grpc (==1.25.0)
-Requires-Dist: pandas (>=2.2.2,<3.0.0)
-Requires-Dist: plotly (>=5.24.1,<6.0.0)
-Requires-Dist: plotly-resampler (>=0.10.0,<0.11.0)
-Requires-Dist: pygooglenews (>=0.1.3,<0.2.0) ; extra == "apps"
+Requires-Dist: pandas (>=2.2.2)
+Requires-Dist: plotly (>=5.24.1)
+Requires-Dist: plotly-resampler (>=0.10.0)
+Requires-Dist: pygooglenews (>=0.1.3) ; extra == "apps"
 Requires-Dist: pylabeladjust (==0.1.13)
-Requires-Dist: pytest (>=8.3.3,<9.0.0) ; extra == "tests"
-Requires-Dist: requests (>=2.32.3,<3.0.0)
-Requires-Dist: scikit-learn (>=1.5.2,<2.0.0)
+Requires-Dist: pytest (>=8.3.3)
+Requires-Dist: requests (>=2.32.3)
+Requires-Dist: scikit-learn (>=1.5.2)
 Requires-Dist: scipy (==1.12)
-Requires-Dist: seaborn (>=0.13.2,<0.14.0)
+Requires-Dist: seaborn (>=0.13.2)
 Requires-Dist: sentence-transformers (==3.0.1)
-Requires-Dist: sentencepiece (>=0.2.0,<0.3.0)
-Requires-Dist: streamlit (>=1.45.0,<2.0.0)
-Requires-Dist: streamlit-code-editor (>=0.1.22,<0.2.0)
-Requires-Dist: thefuzz (>=0.22.1,<0.23.0)
-Requires-Dist: tiktoken (>=0.8.0,<0.9.0)
+Requires-Dist: sentencepiece (>=0.2.0)
+Requires-Dist: starlette (>=0.46.2)
+Requires-Dist: streamlit (>=1.45.0)
+Requires-Dist: streamlit-code-editor (>=0.1.22)
+Requires-Dist: thefuzz (>=0.22.1)
+Requires-Dist: tiktoken (>=0.8.0)
 Requires-Dist: tldextract (>=2.0.1)
 Requires-Dist: tokenizers (<0.20.0)
-Requires-Dist: torch (>=2.5.1,<3.0.0)
-Requires-Dist: tqdm (>=4.67.1,<5.0.0)
-Requires-Dist: umap-learn (>=0.5.7,<0.6.0)
+Requires-Dist: torch (>=2.7.0)
+Requires-Dist: tqdm (>=4.67.1)
+Requires-Dist: umap-learn (>=0.5.7)
+Project-URL: Bug Tracker, https://github.com/rte-france/BERTrend/issues
+Project-URL: Documentation, https://github.com/rte-france/BERTrend
+Project-URL: Homepage, https://github.com/rte-france/BERTrend
 Project-URL: Repository, https://github.com/rte-france/BERTrend
 Description-Content-Type: text/markdown
 
 # BERTrend
 
 [![rte-france - BERTrend](https://img.shields.io/static/v1?label=rte-france&message=BERTrend&color=blue&logo=github)](https://github.com/rte-france/BERTrend "Go to GitHub repo")
-![GitHub License](https://img.shields.io/github/license/rte-france/BERTrend)
 [![stars - BERTrend](https://img.shields.io/github/stars/rte-france/BERTrend?style=social)](https://github.com/rte-france/BERTrend)
 [![forks - BERTrend](https://img.shields.io/github/forks/rte-france/BERTrend?style=social)](https://github.com/rte-france/BERTrend)
 [![issues - BERTrend](https://img.shields.io/github/issues/rte-france/BERTrend)](https://github.com/rte-france/BERTrend/issues)
-![Static Badge](https://img.shields.io/badge/python-%3E%3D3.12-blue)
+[![Downloads](https://pepy.tech/badge/bertrend)](https://pepy.tech/project/bertrend)
+[![PyPi_Version](https://img.shields.io/pypi/v/bertrend.svg)](https://pypi.org/project/bertrend/)
+[![PyPi_Compat](https://img.shields.io/pypi/pyversions/bertrend.svg)](https://pypi.org/project/bertrend/)
+![GitHub License](https://img.shields.io/github/license/rte-france/BERTrend)
 
 ![BERTrend pipeline](./BERTrend_overview.png)
 
-## Work in progress
-
-- make BERTrend available as a python package in PyPI
-
 ## Paper
 
 The code in this repository is part of the work described in the paper: 
@@ -139,7 +137,7 @@ https://github.com/user-attachments/assets/d79368d9-d4e0-4324-8a98-a888f0ab3b65
 - Categorization of signals into three different dataframes: noise, weak and strong signals
   (each dataframe contains useful information about the signals)
 - Interpretation of these signals using LLMs
-- Inspection of topic merging process using Sankey diagram
+- Inspection of the topic merging process using Sankey diagram
 
 ## Installation
 
@@ -164,25 +162,26 @@ BERTrend requires the following environment variables to be set:
 
   `python -m venv <path>/bertrend`
 
+
+- Installation from pypi
+
+  `pip install bertrend` 
+
+
 - Install using pip or poetry in the BERTrend home directory:
 
   ```bash
   # Basic installation with core dependencies
   pip install .  # or poetry install
 
-  # Installation with test dependencies
-  pip install ".[tests]"  # or poetry install --with test
-
-  # Installation with apps dependencies
-  pip install ".[apps]"  # or poetry install --with apps
+    # Installation with apps dependencies
+  pip install ".[apps]"  # or poetry install --extras apps
 
-  # Installation with all optional dependencies
-  pip install ".[tests,apps]"  # or poetry install --with test,apps
   ```
 
 - (Use `pip install -e .` for a local installation)
 
-NB. Due to the number of dependencies, the installation may take some time on an empty environement.
+NB. Due to the number of dependencies, the installation may take some time on an empty environment.
 
 ### Dependency Management
 
@@ -228,7 +227,7 @@ For custom dataset creation, please refer to the documentation page about [data
 
 In order to use the provided demos, datasets :
 - can be stored on the server in: `$BERTREND_BASE_DIR/data/bertrend/`
-- or can be uploaded from the client device on which is displayed the Streamlit apps.
+- or can be uploaded from the client device on which are displayed the Streamlit apps.
 
 #### Dataset Format
 Required columns:
@@ -255,6 +254,13 @@ CUDA_VISIBLE_DEVICES=<gpu_number> streamlit run app.py
 - `Parameters`: Include all BERTopic parameters. Choose the embedding model carefully based on your data language.
 The embedding model can run locally (a GPU is recommended); it is also possible to use a remote embedding service.
 
+We recommend using a remote embedding service (this is now the default config) as it simplifies the computing.
+
+The change of configuration can be done in: [services_default_config.toml](bertrend/config/services_default_config.toml).
+
+An example of a very simple implementation to deploy a separate embedding service can be found here: https://github.com/rte-france/wattelse/tree/main/wattelse/api/embedding. (this implementation is compatible with our code)
+
+
 ![parameters](docs/images/parameters.png)
 
 Please refer to the [BERTopic documentation](https://maartengr.github.io/BERTopic/algorithm/algorithm.html) for parameter details.
@@ -273,7 +279,7 @@ The application is split into several pages:
   - Computes topics over time
   ![exploration_over_time](./docs/images/explore1.png)
 
-  - Allow to select sources and generate short descriptions for each theme
+  - Allow selecting sources and generating short descriptions for each theme
   ![topic_exploration](./docs/images/explore2.png)
 
 - `Topic visualization`: Several visualization types available:
@@ -291,7 +297,7 @@ The application is split into several pages:
   ![temptopic_metrics](./docs/images/temptopic_metrics.png)  
 
 - `Generate Newsletters`
-  - This is an example of application combining topic analysis and LLMs
+  - This is an example of an application combining topic analysis and LLMs
   - Generates newsletters based on prominent topics
   - Uses LLMs to generate topic descriptions and summarize representative documents
 
@@ -340,8 +346,8 @@ Provides various statistics about topics:
 - As topics are hard to interprete by humans, we showcase how it can be
 done using LLMs.
 - Choose a specific topic and click on the button to generate an interpretation
-- LLM-based in-depth signal analysis is done using predefined templates for example:
-  - describe what happens since last timestamp (summary of main events)
+- LLM-based in-depth signal analysis is done using predefined templates, for example:
+  - describe what's happened since the last timestamp (summary of main events)
   - present through a multi-faceted description info about key developments, foreseen changes, potential impact, etc.
 
 
@@ -369,7 +375,7 @@ See [AUTHORS](./AUTHORS.txt) for the full list of contributors.
 
 ## License Information
 
-Copyright 2023-2025 RTE France RTE: http://www.rte-france.com
+Copyright 2023–2025 RTE France RTE: http://www.rte-france.com
 
 This Source Code is subject to the terms of the Mozilla Public License (MPL) v2 also available here: https://www.mozilla.org/en-US/MPL/2.0/
 

{bertrend-0.3.1 → bertrend-0.3.3}/README.md

@@ -1,18 +1,16 @@
 # BERTrend
 
 [![rte-france - BERTrend](https://img.shields.io/static/v1?label=rte-france&message=BERTrend&color=blue&logo=github)](https://github.com/rte-france/BERTrend "Go to GitHub repo")
-![GitHub License](https://img.shields.io/github/license/rte-france/BERTrend)
 [![stars - BERTrend](https://img.shields.io/github/stars/rte-france/BERTrend?style=social)](https://github.com/rte-france/BERTrend)
 [![forks - BERTrend](https://img.shields.io/github/forks/rte-france/BERTrend?style=social)](https://github.com/rte-france/BERTrend)
 [![issues - BERTrend](https://img.shields.io/github/issues/rte-france/BERTrend)](https://github.com/rte-france/BERTrend/issues)
-![Static Badge](https://img.shields.io/badge/python-%3E%3D3.12-blue)
+[![Downloads](https://pepy.tech/badge/bertrend)](https://pepy.tech/project/bertrend)
+[![PyPi_Version](https://img.shields.io/pypi/v/bertrend.svg)](https://pypi.org/project/bertrend/)
+[![PyPi_Compat](https://img.shields.io/pypi/pyversions/bertrend.svg)](https://pypi.org/project/bertrend/)
+![GitHub License](https://img.shields.io/github/license/rte-france/BERTrend)
 
 ![BERTrend pipeline](./BERTrend_overview.png)
 
-## Work in progress
-
-- make BERTrend available as a python package in PyPI
-
 ## Paper
 
 The code in this repository is part of the work described in the paper: 
@@ -67,7 +65,7 @@ https://github.com/user-attachments/assets/d79368d9-d4e0-4324-8a98-a888f0ab3b65
 - Categorization of signals into three different dataframes: noise, weak and strong signals
   (each dataframe contains useful information about the signals)
 - Interpretation of these signals using LLMs
-- Inspection of topic merging process using Sankey diagram
+- Inspection of the topic merging process using Sankey diagram
 
 ## Installation
 
@@ -92,25 +90,26 @@ BERTrend requires the following environment variables to be set:
 
   `python -m venv <path>/bertrend`
 
+
+- Installation from pypi
+
+  `pip install bertrend` 
+
+
 - Install using pip or poetry in the BERTrend home directory:
 
   ```bash
   # Basic installation with core dependencies
   pip install .  # or poetry install
 
-  # Installation with test dependencies
-  pip install ".[tests]"  # or poetry install --with test
-
-  # Installation with apps dependencies
-  pip install ".[apps]"  # or poetry install --with apps
+    # Installation with apps dependencies
+  pip install ".[apps]"  # or poetry install --extras apps
 
-  # Installation with all optional dependencies
-  pip install ".[tests,apps]"  # or poetry install --with test,apps
   ```
 
 - (Use `pip install -e .` for a local installation)
 
-NB. Due to the number of dependencies, the installation may take some time on an empty environement.
+NB. Due to the number of dependencies, the installation may take some time on an empty environment.
 
 ### Dependency Management
 
@@ -156,7 +155,7 @@ For custom dataset creation, please refer to the documentation page about [data
 
 In order to use the provided demos, datasets :
 - can be stored on the server in: `$BERTREND_BASE_DIR/data/bertrend/`
-- or can be uploaded from the client device on which is displayed the Streamlit apps.
+- or can be uploaded from the client device on which are displayed the Streamlit apps.
 
 #### Dataset Format
 Required columns:
@@ -183,6 +182,13 @@ CUDA_VISIBLE_DEVICES=<gpu_number> streamlit run app.py
 - `Parameters`: Include all BERTopic parameters. Choose the embedding model carefully based on your data language.
 The embedding model can run locally (a GPU is recommended); it is also possible to use a remote embedding service.
 
+We recommend using a remote embedding service (this is now the default config) as it simplifies the computing.
+
+The change of configuration can be done in: [services_default_config.toml](bertrend/config/services_default_config.toml).
+
+An example of a very simple implementation to deploy a separate embedding service can be found here: https://github.com/rte-france/wattelse/tree/main/wattelse/api/embedding. (this implementation is compatible with our code)
+
+
 ![parameters](docs/images/parameters.png)
 
 Please refer to the [BERTopic documentation](https://maartengr.github.io/BERTopic/algorithm/algorithm.html) for parameter details.
@@ -201,7 +207,7 @@ The application is split into several pages:
   - Computes topics over time
   ![exploration_over_time](./docs/images/explore1.png)
 
-  - Allow to select sources and generate short descriptions for each theme
+  - Allow selecting sources and generating short descriptions for each theme
   ![topic_exploration](./docs/images/explore2.png)
 
 - `Topic visualization`: Several visualization types available:
@@ -219,7 +225,7 @@ The application is split into several pages:
   ![temptopic_metrics](./docs/images/temptopic_metrics.png)  
 
 - `Generate Newsletters`
-  - This is an example of application combining topic analysis and LLMs
+  - This is an example of an application combining topic analysis and LLMs
   - Generates newsletters based on prominent topics
   - Uses LLMs to generate topic descriptions and summarize representative documents
 
@@ -268,8 +274,8 @@ Provides various statistics about topics:
 - As topics are hard to interprete by humans, we showcase how it can be
 done using LLMs.
 - Choose a specific topic and click on the button to generate an interpretation
-- LLM-based in-depth signal analysis is done using predefined templates for example:
-  - describe what happens since last timestamp (summary of main events)
+- LLM-based in-depth signal analysis is done using predefined templates, for example:
+  - describe what's happened since the last timestamp (summary of main events)
   - present through a multi-faceted description info about key developments, foreseen changes, potential impact, etc.
 
 
@@ -297,7 +303,7 @@ See [AUTHORS](./AUTHORS.txt) for the full list of contributors.
 
 ## License Information
 
-Copyright 2023-2025 RTE France RTE: http://www.rte-france.com
+Copyright 2023–2025 RTE France RTE: http://www.rte-france.com
 
 This Source Code is subject to the terms of the Mozilla Public License (MPL) v2 also available here: https://www.mozilla.org/en-US/MPL/2.0/
 

{bertrend-0.3.1 → bertrend-0.3.3}/bertrend/llm_utils/newsletter_features.py

@@ -126,9 +126,9 @@ def generate_newsletter(
     # Iterate over topics
     for i in tqdm(range(top_n_topics), desc="Processing topics..."):
         sub_df = get_most_representative_docs(
-            topic_model,
-            df,
-            topics,
+            topic_model=topic_model,
+            df=df,
+            topics=topics,
             mode=top_n_docs_mode,
             df_split=df_split,
             topic_number=i,

{bertrend-0.3.1 → bertrend-0.3.3}/bertrend/services/authentication.py

@@ -2,7 +2,6 @@
 #  See AUTHORS.txt
 #  SPDX-License-Identifier: MPL-2.0
 #  This file is part of BERTrend.
-import os
 import time
 
 import requests

{bertrend-0.3.1 → bertrend-0.3.3}/bertrend/services/summary/abstractive_summarizer.py

@@ -22,7 +22,7 @@ class AbstractiveSummarizer(Summarizer):
     def __init__(self, model_name=DEFAULT_ABSTRACTIVE_MODEL):
         self.model_name = model_name
         self.WHITESPACE_HANDLER = lambda k: re.sub(
-            "\s+", " ", re.sub("\n+", " ", k.strip())
+            r"\s+", " ", re.sub(r"\n+", " ", k.strip())
         )
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

{bertrend-0.3.1 → bertrend-0.3.3}/bertrend/services/summary/chatgpt_summarizer.py

@@ -35,7 +35,7 @@ class GPTSummarizer(Summarizer):
         prompt_language: str = "fr",
         max_article_length: int = 1500,
         model_name: str = None,
-        **kwargs
+        **kwargs,
     ) -> str:
         # Limit input length in case the text is large
         article_text = keep_first_n_words(article_text, max_article_length)

bertrend-0.3.3/bertrend/tests/README_integration_tests.md

@@ -0,0 +1,91 @@
+# BERTrend Integration Tests
+
+This directory contains integration tests for the BERTrend project. These tests verify that the different components of the system work together correctly in real-world scenarios.
+
+## Overview
+
+The integration tests are organized into several test classes, each focusing on a specific integration point:
+
+1. **TestEmbeddingServiceIntegration**: Tests the integration of the EmbeddingService with local embedding models.
+2. **TestBERTopicModelIntegration**: Tests the integration of BERTopicModel with pre-computed embeddings.
+3. **TestBERTrendIntegration**: Tests the end-to-end workflow of BERTrend, including training topic models, calculating signal popularity, and classifying signals.
+4. **TestDataPipelineIntegration**: Tests the data loading and processing pipeline.
+5. **TestNewsletterIntegration**: Tests the newsletter generation functionality, which combines topic modeling with LLM-based summarization.
+
+## Running the Tests
+
+To run the integration tests, you need to have the BERTrend package installed with all its dependencies. You can run the tests using pytest with the following command:
+
+```bash
+pytest -m integration bertrend/tests/test_integration.py -v
+```
+
+This will run only the tests marked with the `integration` marker. The `-v` flag provides verbose output.
+
+### Running Specific Tests
+
+To run a specific test class, you can use:
+
+```bash
+pytest -m integration bertrend/tests/test_integration.py::TestBERTrendIntegration -v
+```
+
+To run a specific test method, you can use:
+
+```bash
+pytest -m integration bertrend/tests/test_integration.py::TestBERTrendIntegration::test_bertrend_end_to_end -v
+```
+
+## Test Requirements
+
+The integration tests have the following requirements:
+
+1. **Sentence Transformer Models**: The tests use the `all-MiniLM-L6-v2` model, which will be downloaded automatically if not already present.
+2. **OpenAI API Key**: The newsletter generation test requires an OpenAI API key. If the key is not available, the test will be skipped.
+3. **Temporary Directory**: The tests use pytest's `tmp_path` fixture to create temporary directories for saving models and files.
+
+## Adding New Integration Tests
+
+When adding new integration tests, please follow these guidelines:
+
+1. Use the `@pytest.mark.integration` decorator to mark the test as an integration test.
+2. Use small, fast models and synthetic data to ensure the tests run efficiently.
+3. Include assertions that verify the correct integration between components.
+4. Handle external dependencies gracefully, skipping tests if necessary.
+5. Clean up any resources created during the test.
+
+## Example
+
+Here's an example of a simple integration test:
+
+```python
+@pytest.mark.integration
+def test_embedding_service_with_local_model():
+    """Test that EmbeddingService works with a local model."""
+    # Use a small, fast model for testing
+    model_name = "all-MiniLM-L6-v2"
+    
+    # Create a small test dataset
+    texts = [
+        "This is a test document about artificial intelligence.",
+        "Machine learning is a subset of artificial intelligence.",
+        "Natural language processing is used in many applications."
+    ]
+    
+    # Initialize the embedding service with local model
+    embedding_service = EmbeddingService(
+        local=True,
+        model_name=model_name
+    )
+    
+    # Generate embeddings
+    embeddings, token_strings, token_embeddings = embedding_service.embed(texts)
+    
+    # Verify the results
+    assert embeddings is not None
+    assert embeddings.shape == (len(texts), 384)  # 384 is the dimension for all-MiniLM-L6-v2
+    assert token_strings is not None
+    assert token_embeddings is not None
+    assert len(token_strings) == len(texts)
+    assert len(token_embeddings) == len(texts)
+```

bertrend-0.3.3/bertrend/tests/conftest.py

@@ -0,0 +1,11 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+import pytest
+
+
+def pytest_configure(config):
+    """Register custom markers."""
+    config.addinivalue_line("markers", "integration: mark test as an integration test")

{bertrend-0.3.1 → bertrend-0.3.3}/bertrend/tests/test_abstractive_summarizer.py

@@ -36,12 +36,15 @@ def mock_model():
 @pytest.fixture
 def mock_abstractive_summarizer(mock_tokenizer, mock_model):
     """Create a mock AbstractiveSummarizer with mocked dependencies."""
-    with patch(
-        "bertrend.services.summary.abstractive_summarizer.AutoTokenizer.from_pretrained",
-        return_value=mock_tokenizer,
-    ), patch(
-        "bertrend.services.summary.abstractive_summarizer.AutoModelForSeq2SeqLM.from_pretrained",
-        return_value=mock_model,
+    with (
+        patch(
+            "bertrend.services.summary.abstractive_summarizer.AutoTokenizer.from_pretrained",
+            return_value=mock_tokenizer,
+        ),
+        patch(
+            "bertrend.services.summary.abstractive_summarizer.AutoModelForSeq2SeqLM.from_pretrained",
+            return_value=mock_model,
+        ),
     ):
         summarizer = AbstractiveSummarizer()
         # Set the mocks directly to ensure they're used
@@ -52,11 +55,14 @@ def mock_abstractive_summarizer(mock_tokenizer, mock_model):
 
 def test_initialization():
     """Test that the AbstractiveSummarizer initializes with the correct model."""
-    with patch(
-        "bertrend.services.summary.abstractive_summarizer.AutoTokenizer.from_pretrained"
-    ) as mock_tokenizer, patch(
-        "bertrend.services.summary.abstractive_summarizer.AutoModelForSeq2SeqLM.from_pretrained"
-    ) as mock_model:
+    with (
+        patch(
+            "bertrend.services.summary.abstractive_summarizer.AutoTokenizer.from_pretrained"
+        ) as mock_tokenizer,
+        patch(
+            "bertrend.services.summary.abstractive_summarizer.AutoModelForSeq2SeqLM.from_pretrained"
+        ) as mock_model,
+    ):
 
         # Mock the to() method to avoid device-related issues
         mock_model.return_value.to.return_value = mock_model.return_value
@@ -124,11 +130,14 @@ def test_custom_model():
     """Test initialization with a custom model."""
     custom_model = "custom/model"
 
-    with patch(
-        "bertrend.services.summary.abstractive_summarizer.AutoTokenizer.from_pretrained"
-    ) as mock_tokenizer, patch(
-        "bertrend.services.summary.abstractive_summarizer.AutoModelForSeq2SeqLM.from_pretrained"
-    ) as mock_model:
+    with (
+        patch(
+            "bertrend.services.summary.abstractive_summarizer.AutoTokenizer.from_pretrained"
+        ) as mock_tokenizer,
+        patch(
+            "bertrend.services.summary.abstractive_summarizer.AutoModelForSeq2SeqLM.from_pretrained"
+        ) as mock_model,
+    ):
 
         # Mock the to() method to avoid device-related issues
         mock_model.return_value.to.return_value = mock_model.return_value