bertrend 0.3.4__tar.gz → 0.3.5__tar.gz

{bertrend-0.3.4 → bertrend-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: bertrend
-Version: 0.3.4
+Version: 0.3.5
 Summary: BERTrend analyses topic evolution over time using state-of-the-art transformer models: topic modelling and LLM
 License: LICENSE.md
 Keywords: topics,BERTopic,transformers,weak signals,LLM
@@ -11,6 +11,7 @@ Maintainer-email: jerome.picault@rte-france.com
 Requires-Python: >=3.12,<4.0
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Provides-Extra: apps
+Requires-Dist: PyJWT (>=2.10.1)
 Requires-Dist: Pyqtree (==1.0.0)
 Requires-Dist: arxiv (>=2.2.0) ; extra == "apps"
 Requires-Dist: bertopic (==0.16.2)
@@ -48,7 +49,7 @@ 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)
+Requires-Dist: pytest (>=8.4.0)
 Requires-Dist: requests (>=2.32.3)
 Requires-Dist: scikit-learn (>=1.5.2)
 Requires-Dist: scipy (==1.12)
@@ -56,7 +57,7 @@ Requires-Dist: seaborn (>=0.13.2)
 Requires-Dist: sentence-transformers (==3.0.1)
 Requires-Dist: sentencepiece (>=0.2.0)
 Requires-Dist: starlette (>=0.46.2)
-Requires-Dist: streamlit (>=1.45.0)
+Requires-Dist: streamlit (>=1.45.1)
 Requires-Dist: streamlit-code-editor (>=0.1.22)
 Requires-Dist: thefuzz (>=0.22.1)
 Requires-Dist: tiktoken (>=0.8.0)
@@ -259,7 +260,7 @@ We recommend using a remote embedding service (this is now the default config) a
 
 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)
+An example of a very simple implementation to deploy a separate embedding service can be found here: `bertrend/services/embedding_server`
 
 
 ![parameters](docs/images/parameters.png)

{bertrend-0.3.4 → bertrend-0.3.5}/README.md

@@ -186,7 +186,7 @@ We recommend using a remote embedding service (this is now the default config) a
 
 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)
+An example of a very simple implementation to deploy a separate embedding service can be found here: `bertrend/services/embedding_server`
 
 
 ![parameters](docs/images/parameters.png)

bertrend-0.3.5/bertrend/services/embedding_server/README.md

@@ -0,0 +1,31 @@
+# How to use the embedding server
+
+## Launch the embedding server
+
+- check the file `config/default_config.toml`
+
+- run the embedding server
+
+```
+python start.py
+```
+
+## Check the embedding server
+Ensure the API is launched.
+
+```python
+from bertrend.services.embedding_client import EmbeddingAPIClient
+
+embedding_service_endpoint = "https://yourservice:1234"
+
+# Initialize the API
+api = EmbeddingAPIClient(embedding_service_endpoint)
+
+# Get embedding model name
+print(api.get_api_model_name())
+
+# Transform text to embedding
+text = "Hello, world!"
+embedding = api.embed_query(text)
+print(embedding)
+```

bertrend-0.3.5/bertrend/services/embedding_server/config/default_config.toml

@@ -0,0 +1,7 @@
+host = "0.0.0.0"
+port = 6464
+model_name = "OrdalieTech/Solon-embeddings-large-0.1"
+#model_name=dangvantuan/sentence-camembert-large
+#model_name=antoinelouis/biencoder-camembert-base-mmarcoFR
+number_workers = 2
+cuda_visible_devices = "0"

bertrend-0.3.5/bertrend/services/embedding_server/config/settings.py

@@ -0,0 +1,34 @@
+#  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 os
+from pathlib import Path
+import tomllib
+from pydantic_settings import BaseSettings
+
+DEFAULT_CONFIG_FILE = Path(__file__).parent / "default_config.toml"
+
+
+class EmbeddingAPIConfig(BaseSettings):
+    host: str
+    port: int
+    model_name: str
+    number_workers: int
+    cuda_visible_devices: str
+
+
+# Use environment variable to specify which config to load
+def get_config() -> EmbeddingAPIConfig:
+    """
+    Return the configuration for the embedding API.
+    The configuration is loaded from a toml file.
+    The path to the toml file can be specified using the environment variable `EMBEDDING_API_CONFIG_FILE`.
+    If the environment variable is not set, the default configuration file is used.
+    """
+    config_file = os.environ.get("EMBEDDING_API_CONFIG_FILE", None)
+    if config_file is None:
+        config_file = DEFAULT_CONFIG_FILE
+    with open(config_file, "rb") as f:
+        return EmbeddingAPIConfig(**tomllib.load(f))

bertrend-0.3.5/bertrend/services/embedding_server/main.py

@@ -0,0 +1,14 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+from fastapi import FastAPI
+
+from bertrend.services.embedding_server.routers import embeddings, info, authentication
+
+app = FastAPI()
+
+app.include_router(info.router, tags=["Info"])
+app.include_router(embeddings.router, tags=["Embedding"])
+app.include_router(authentication.router, tags=["Authentication"])

bertrend-0.3.5/bertrend/services/embedding_server/models.py

@@ -0,0 +1,12 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+# Pydantic class for FastAPI typing control
+from pydantic import BaseModel
+
+
+class InputText(BaseModel):
+    text: str | list[str]
+    show_progress_bar: bool

bertrend-0.3.5/bertrend/services/embedding_server/routers/authentication.py

@@ -0,0 +1,45 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+from typing import Annotated
+
+from fastapi import Depends, APIRouter, Security
+from fastapi.security import OAuth2PasswordRequestForm
+
+from bertrend.services.embedding_server.security import (
+    ADMIN,
+    Token,
+    TokenData,
+    get_token,
+    list_registered_clients,
+    get_current_client,
+    view_rate_limits,
+)
+
+router = APIRouter()
+
+
+@router.post("/token")
+async def login_for_access_token(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+) -> Token:
+    return get_token(form_data)
+
+
+@router.get(
+    "/list_registered_clients", summary="List registered clients (requires admin scope)"
+)
+async def list_clients(
+    current_client: TokenData = Security(get_current_client, scopes=[ADMIN])
+):
+    return list_registered_clients()
+
+
+@router.get(
+    "/rate-limits", summary="Check current rate limit usage (requires admin scope)"
+)
+async def get_rate_limits(
+    current_client: TokenData = Security(get_current_client, scopes=[ADMIN]),
+):
+    return view_rate_limits()

bertrend-0.3.5/bertrend/services/embedding_server/routers/embeddings.py

@@ -0,0 +1,39 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+from fastapi import APIRouter, Security
+from loguru import logger
+from sentence_transformers import SentenceTransformer
+
+from bertrend.services.embedding_server.security import (
+    TokenData,
+    get_current_client,
+    FULL_ACCESS,
+)
+
+from bertrend.services.embedding_server.config.settings import get_config
+from bertrend.services.embedding_server.models import InputText
+
+# Load the configuration
+CONFIG = get_config()
+
+# Load embedding model
+logger.debug(f"Loading embedding model : {CONFIG.model_name}")
+EMBEDDING_MODEL = SentenceTransformer(CONFIG.model_name, trust_remote_code=True)
+
+# Fix max model length error
+if EMBEDDING_MODEL.max_seq_length == 514:
+    EMBEDDING_MODEL.max_seq_length = 512
+router = APIRouter()
+
+
+@router.post("/encode", summary="Embed data (requires full_access scope)")
+def embed(
+    input: InputText,
+    current_client: TokenData = Security(get_current_client, scopes=[FULL_ACCESS]),
+):
+    logger.debug(f"Request by: {current_client.client_id}")
+    emb = EMBEDDING_MODEL.encode(input.text, show_progress_bar=input.show_progress_bar)
+    return {"embeddings": emb.tolist()}

bertrend-0.3.5/bertrend/services/embedding_server/routers/info.py

@@ -0,0 +1,29 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+from fastapi import APIRouter
+
+from bertrend.services.embedding_server.config.settings import get_config
+
+# Load the configuration
+CONFIG = get_config()
+
+
+router = APIRouter()
+
+
+@router.get("/health")
+def health():
+    return {"status": "ok"}
+
+
+@router.get("/model_name")
+def get_model_name():
+    return CONFIG.model_name
+
+
+@router.get("/num_workers")
+def get_num_workers():
+    return CONFIG.number_workers

bertrend-0.3.5/bertrend/services/embedding_server/security.py

@@ -0,0 +1,368 @@
+import json
+import secrets
+import time
+from collections import defaultdict, deque
+
+import urllib3
+from fastapi import Depends, HTTPException, status, Request
+from fastapi.security import (
+    OAuth2PasswordRequestForm,
+    SecurityScopes,
+)
+from fastapi.security.oauth2 import OAuth2PasswordBearer
+import jwt  # PyJWT
+from datetime import datetime, timedelta, timezone
+from typing import Annotated, Deque
+
+from jwt import InvalidTokenError
+from loguru import logger
+from pydantic import BaseModel, ValidationError
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Disable warnings related to https certificates
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+
+def generate_hex_token(length: int = 32):
+    """Generates a random hexadecimal string of the specified length."""
+    return secrets.token_hex(length)
+
+
+# Configuration (should be in environment variables)
+# to get a string like this run:
+# openssl rand -hex 32
+SECRET_KEY = os.getenv("BERTREND_SECRET_KEY", generate_hex_token())
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 48  # 48h
+
+# Rate limiting configuration
+DEFAULT_RATE_LIMIT = int(os.getenv("DEFAULT_RATE_LIMIT", "50"))  # Requests per minute
+DEFAULT_RATE_WINDOW = int(os.getenv("DEFAULT_RATE_WINDOW", "60"))  # Window in seconds
+
+# Client registry - in production, store somewhere else - database, file...
+CLIENT_REGISTRY_FILE = os.getenv(
+    "CLIENT_REGISTRY_FILE", f"{os.path.expanduser('~')}/bertrend_client_registry.json"
+)
+
+RESTRICTED_ACCESS = "restricted_access"
+FULL_ACCESS = "full_access"
+ADMIN = "admin"
+
+# Define available scopes with descriptions
+SCOPES = {
+    RESTRICTED_ACCESS: "Access limited to some endpoints",
+    FULL_ACCESS: "Full access to endpoints except admin endpoints",
+    ADMIN: "Admin access",
+}
+
+
+DEFAULT_CLIENT_REGISTRY = {
+    "admin": {
+        "client_secret": generate_hex_token(),
+        "scopes": [ADMIN, FULL_ACCESS, RESTRICTED_ACCESS],
+        "authorized_groups": [],  # Empty list means access to all groups
+        "rate_limit": DEFAULT_RATE_LIMIT * 2,  # Higher limit for admin
+        "rate_window": DEFAULT_RATE_WINDOW,
+    },
+    "bertrend": {
+        "client_secret": generate_hex_token(),
+        "scopes": [RESTRICTED_ACCESS],
+        "authorized_groups": [],  # Empty list means access to all groups
+        "rate_limit": DEFAULT_RATE_LIMIT,
+        "rate_window": DEFAULT_RATE_WINDOW,
+    },
+}
+
+# Configure OAuth2
+oauth2_scheme = OAuth2PasswordBearer(
+    tokenUrl="token",
+    scopes=SCOPES,
+)
+
+# Store for rate limiting (in memory - for production use Redis or similar)
+# Maps client_id -> deque of request timestamps
+rate_limit_store: dict[str, Deque[float]] = defaultdict(lambda: deque(maxlen=1000))
+
+
+def load_client_registry():
+    """Function to load client registry from JSON file"""
+    try:
+        with open(CLIENT_REGISTRY_FILE, "r") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        # Create a default client registry file if it doesn't exist
+        default_registry = DEFAULT_CLIENT_REGISTRY
+        with open(CLIENT_REGISTRY_FILE, "w") as f:
+            json.dump(default_registry, f, indent=4)
+        return default_registry
+    except json.JSONDecodeError:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Client registry file is not valid JSON",
+        )
+
+
+def list_registered_clients():
+    """View registered clients (without exposing secrets)"""
+    client_registry = load_client_registry()
+    # Remove sensitive information like client secrets
+    clients_info = {}
+    for client_id, client_data in client_registry.items():
+        clients_info[client_id] = {
+            "scopes": client_data["scopes"],
+            "authorized_groups": client_data.get("authorized_groups", []),
+            "rate_limit": client_data.get("rate_limit", DEFAULT_RATE_LIMIT),
+            "rate_window": client_data.get("rate_window", DEFAULT_RATE_WINDOW),
+        }
+    return clients_info
+
+
+def view_rate_limits():
+    """Return the current usage of the API per client"""
+    current_time = time.time()
+    usage = {}
+
+    for client_id, requests in rate_limit_store.items():
+        # Count only requests within the window
+        client_registry = load_client_registry()
+        if client_id in client_registry:
+            window = client_registry[client_id].get("rate_window", DEFAULT_RATE_WINDOW)
+            cutoff_time = current_time - window
+            active_requests = sum(1 for t in requests if t >= cutoff_time)
+
+            usage[client_id] = {
+                "usage": active_requests,
+                "limit": client_registry[client_id].get(
+                    "rate_limit", DEFAULT_RATE_LIMIT
+                ),
+                "window": window,
+            }
+
+    return usage
+
+
+class Token(BaseModel):
+    access_token: str
+    token_type: str
+    expires_in: float  # This represents the token expiration time in seconds
+
+
+class TokenData(BaseModel):
+    client_id: str | None = None
+    scopes: list[str] = []
+    authorized_groups: list[str] = []
+    rate_limit: int = DEFAULT_RATE_LIMIT
+    rate_window: int = DEFAULT_RATE_WINDOW
+
+
+def create_access_token(data: dict, expires_delta: timedelta | None = None):
+    """Creates access token"""
+    to_encode = data.copy()
+    expire = datetime.now(timezone.utc) + expires_delta or timedelta(
+        minutes=ACCESS_TOKEN_EXPIRE_MINUTES
+    )
+    to_encode.update(
+        {
+            "exp": expire,
+            "expires_in": (
+                expires_delta.total_seconds()
+                if expires_delta
+                else ACCESS_TOKEN_EXPIRE_MINUTES * 60
+            ),
+        }
+    )
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+
+async def check_rate_limit(request: Request, token_data: TokenData):
+    """
+    Check if the client has exceeded their rate limit
+
+    Args:
+        request: The FastAPI request object
+        token_data: The TokenData object containing client rate limit info
+
+    Raises:
+        HTTPException: If rate limit is exceeded
+    """
+    if token_data.client_id is None:
+        return
+
+    client_id = token_data.client_id
+    rate_limit = token_data.rate_limit
+    rate_window = token_data.rate_window
+
+    # Record this request
+    current_time = time.time()
+    client_requests = rate_limit_store[client_id]
+    client_requests.append(current_time)
+
+    # Clean up old requests outside the window
+    cutoff_time = current_time - rate_window
+    while client_requests and client_requests[0] < cutoff_time:
+        client_requests.popleft()
+
+    # Check if rate limit exceeded
+    if len(client_requests) > rate_limit:
+        # Calculate reset time
+        reset_time = int(client_requests[0] + rate_window - current_time)
+        raise HTTPException(
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+            detail=f"Rate limit exceeded. Try again in {reset_time} seconds.",
+            headers={"Retry-After": str(reset_time), "X-Rate-Limit": str(rate_limit)},
+        )
+
+
+async def get_current_client(
+    security_scopes: SecurityScopes,
+    token: Annotated[str, Depends(oauth2_scheme)],
+    request: Request,
+):
+    """Function to verify token and extract client data"""
+    if security_scopes.scopes:
+        authenticate_value = f'Bearer scope="{security_scopes.scope_str}"'
+    else:
+        authenticate_value = "Bearer"
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": authenticate_value},
+    )
+    try:
+        payload = jwt.decode(token, key=SECRET_KEY, algorithms=[ALGORITHM])
+        client_id = payload.get("sub")
+        if client_id is None:
+            raise credentials_exception
+        token_scopes = payload.get("scopes", [])
+        token_groups = payload.get("authorized_groups", [])
+        rate_limit = payload.get("rate_limit", DEFAULT_RATE_LIMIT)
+        rate_window = payload.get("rate_window", DEFAULT_RATE_WINDOW)
+
+        token_data = TokenData(
+            client_id=client_id,
+            scopes=token_scopes,
+            authorized_groups=token_groups,
+            rate_limit=rate_limit,
+            rate_window=rate_window,
+        )
+    except (InvalidTokenError, ValidationError):
+        raise credentials_exception
+
+    # Check if client still exists
+    if client_id not in load_client_registry():
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Client no longer exists",
+            headers={"WWW-Authenticate": authenticate_value},
+        )
+
+    # Check if client has necessary scopes
+    for scope in security_scopes.scopes:
+        if scope not in token_data.scopes:
+            logger.error(f"{scope} not in {token_data.scopes}, not enough permissions")
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail=f"Not enough permissions. Required scope: {scope}",
+                headers={"WWW-Authenticate": authenticate_value},
+            )
+
+    # Check rate limit
+    await check_rate_limit(request, token_data)
+
+    return token_data
+
+
+def get_token(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+) -> Token:
+    """Implement the token endpoint for token generation"""
+    client_id = (
+        form_data.username
+    )  # In client credentials, username field holds the client_id
+    client_secret = form_data.password  # Password field holds the client_secret
+
+    # Load client registry
+    client_registry = load_client_registry()
+
+    # Validate client
+    if client_id not in client_registry:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid client credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    client = client_registry[client_id]
+    if client["client_secret"] != client_secret:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid client credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    # Validate scopes
+    client_scopes = client["scopes"]
+    requested_scopes = form_data.scopes
+
+    # If no scopes requested, grant all available to client
+    if not requested_scopes:
+        scopes = client_scopes
+    else:
+        # Check if requested scopes are valid
+        for scope in requested_scopes:
+            if scope not in client_scopes:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail=f"Client '{client_id}' does not have access to scope '{scope}'",
+                    headers={"WWW-Authenticate": "Bearer"},
+                )
+        scopes = requested_scopes
+
+    # Get authorized groups and rate limit
+    authorized_groups = client.get("authorized_groups", [])
+    rate_limit = client.get("rate_limit", DEFAULT_RATE_LIMIT)
+    rate_window = client.get("rate_window", DEFAULT_RATE_WINDOW)
+
+    # Create access token
+    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    access_token = create_access_token(
+        data={
+            "sub": client_id,
+            "scopes": scopes,
+            "authorized_groups": authorized_groups,
+            "rate_limit": rate_limit,
+            "rate_window": rate_window,
+        },
+        expires_delta=access_token_expires,
+    )
+
+    expires_in = (
+        access_token_expires.total_seconds()
+    )  # Calculate the expiration time in seconds
+
+    logger.debug(f"Token created for client '{client_id}'")
+    return Token(access_token=access_token, token_type="bearer", expires_in=expires_in)
+
+
+def is_authorized_for_group(token_data: TokenData, group: str) -> bool:
+    """
+    Check if a client is authorized for a specific group
+
+    Args:
+        token_data: The TokenData object containing client information
+        group: The group to check authorization for
+
+    Returns:
+        bool: True if authorized, False otherwise
+    """
+    # Empty authorized_groups means access to all groups
+    if not token_data.authorized_groups:
+        return True
+
+    # Otherwise, check if the requested group is in the authorized groups
+    return group in token_data.authorized_groups

bertrend-0.3.5/bertrend/services/embedding_server/start.py

@@ -0,0 +1,26 @@
+#  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 os
+import uvicorn
+
+from bertrend.services.embedding_server.config.settings import get_config
+
+# Load the configuration
+CONFIG = get_config()
+
+# Set the CUDA_VISIBLE_DEVICES environment variable
+os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG.cuda_visible_devices
+
+# Start the FastAPI application
+if __name__ == "__main__":
+    uvicorn.run(
+        "wattelse.api.embedding.main:app",
+        host=CONFIG.host,
+        port=CONFIG.port,
+        workers=CONFIG.number_workers,
+        ssl_keyfile="../key.pem",
+        ssl_certfile="../cert.pem",
+    )