Merge branch 'main' into websocket-unit-tests

Author: kevinmessiaen

giskard/commands/cli_worker.py

@@ -136,6 +136,7 @@ def initialize_hf_token(hf_token, is_server):
 def _start_command(is_server, url: AnyHttpUrl, api_key, is_daemon, hf_token=None, nb_workers=None):
     from giskard.ml_worker.ml_worker import MLWorker
 
+    os.environ["TQDM_DISABLE"] = "1"
     start_msg = "Starting ML Worker"
     start_msg += " server" if is_server else " client"
     if is_daemon:

giskard/ml_worker/ml_worker.py

@@ -161,6 +161,7 @@ async def start(self, nb_workers: Optional[int] = None):
                 # as described in https://github.com/jasonrbriggs/stomp.py/issues/424
                 # and https://github.com/websocket-client/websocket-client/issues/930
                 logger.warn(f"WebSocket connection may not be properly closed: {e}")
+                logger.exception(e)
 
     def stop(self):
         if self.ws_conn:

giskard/ml_worker/websocket/listener.py

@@ -78,17 +78,25 @@ def websocket_log_actor(ml_worker: MLWorkerInfo, req: Dict, *args, **kwargs):
 WEBSOCKET_ACTORS = dict((action.name, websocket_log_actor) for action in MLWorkerAction)
 
 
-def wrapped_handle_result(action: MLWorkerAction, ml_worker: MLWorker, start: float, rep_id: Optional[str]):
+def wrapped_handle_result(
+    action: MLWorkerAction, ml_worker: MLWorker, start: float, rep_id: Optional[str], ignore_timeout: bool
+):
     def handle_result(future: Union[Future, Callable[..., websocket.WorkerReply]]):
         log_pool_stats()
 
         info = None  # Needs to be defined in case of cancellation
 
         try:
             info: websocket.WorkerReply = future.result() if isinstance(future, Future) else future()
-        except CancelledError:
-            info: websocket.WorkerReply = websocket.Empty()
-            logger.warning("Task for %s has timed out and been cancelled", action.name)
+        except CancelledError as e:
+            if ignore_timeout:
+                info: websocket.WorkerReply = websocket.Empty()
+                logger.warning("Task for %s has timed out and been cancelled", action.name)
+            else:
+                info: websocket.WorkerReply = websocket.ErrorReply(
+                    error_str=str(e), error_type=type(e).__name__, detail=traceback.format_exc()
+                )
+                logger.warning(e)
         except Exception as e:
             info: websocket.WorkerReply = websocket.ErrorReply(
                 error_str=str(e), error_type=type(e).__name__, detail=traceback.format_exc()
@@ -171,7 +179,7 @@ def parse_and_execute(
     )
 
 
-def dispatch_action(callback, ml_worker, action, req, execute_in_pool, timeout=None):
+def dispatch_action(callback, ml_worker, action, req, execute_in_pool, timeout=None, ignore_timeout=False):
     # Parse the response ID
     rep_id = req["id"] if "id" in req.keys() else None
     # Parse the param
@@ -199,7 +207,7 @@ def dispatch_action(callback, ml_worker, action, req, execute_in_pool, timeout=N
     )
     start = time.process_time()
 
-    result_handler = wrapped_handle_result(action, ml_worker, start, rep_id)
+    result_handler = wrapped_handle_result(action, ml_worker, start, rep_id, ignore_timeout=ignore_timeout)
     # If execution should be done in a pool
     if execute_in_pool:
         logger.debug("Submitting for action %s '%s' into the pool", action.name, callback.__name__)
@@ -227,7 +235,9 @@ def dispatch_action(callback, ml_worker, action, req, execute_in_pool, timeout=N
     )
 
 
-def websocket_actor(action: MLWorkerAction, execute_in_pool: bool = True, timeout: Optional[float] = None):
+def websocket_actor(
+    action: MLWorkerAction, execute_in_pool: bool = True, timeout: Optional[float] = None, ignore_timeout: bool = False
+):
     """
     Register a function as an actor to an action from WebSocket connection
     """
@@ -238,7 +248,7 @@ def websocket_actor_callback(callback: callable):
         logger.debug(f'Registered "{callback.__name__}" for ML Worker "{action.name}"')
 
         def wrapped_callback(ml_worker: MLWorker, req: dict, *args, **kwargs):
-            dispatch_action(callback, ml_worker, action, req, execute_in_pool, timeout)
+            dispatch_action(callback, ml_worker, action, req, execute_in_pool, timeout, ignore_timeout)
 
         WEBSOCKET_ACTORS[action.name] = wrapped_callback
 
@@ -664,7 +674,7 @@ def echo(params: websocket.EchoMsg, *args, **kwargs) -> websocket.EchoMsg:
     return params
 
 
-@websocket_actor(MLWorkerAction.getPush, timeout=30)
+@websocket_actor(MLWorkerAction.getPush, timeout=30, ignore_timeout=True)
 def get_push(
     client: Optional[GiskardClient], params: websocket.GetPushParam, *args, **kwargs
 ) -> websocket.GetPushResponse:

giskard/models/model_explanation.py

@@ -1,14 +1,15 @@
+from typing import Any, Callable, Dict, List
+
 import logging
 import warnings
-from typing import Callable, Dict, List, Any
 
 import numpy as np
 import pandas as pd
 
+from giskard.core.errors import GiskardImportError
 from giskard.datasets.base import Dataset
-from giskard.models.base import BaseModel
 from giskard.ml_worker.utils.logging import timer
-from giskard.core.errors import GiskardImportError
+from giskard.models.base import BaseModel
 
 warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")
 logger = logging.getLogger(__name__)
@@ -189,6 +190,7 @@ def explain_with_shap(model: BaseModel, dataset: Dataset, only_highest_proba: bo
 
     try:
         from shap import Explanation
+
         from giskard.models.shap_result import ShapResult
     except ImportError as e:
         raise GiskardImportError("shap") from e
@@ -244,8 +246,8 @@ def explain(model: BaseModel, dataset: Dataset, input_data: Dict):
 @timer()
 def explain_text(model: BaseModel, input_df: pd.DataFrame, text_column: str, text_document: str):
     try:
-        from shap.maskers import Text
         from shap import Explainer
+        from shap.maskers import Text
     except ImportError as e:
         raise GiskardImportError("shap") from e
     try:
@@ -259,7 +261,8 @@ def explain_text(model: BaseModel, input_df: pd.DataFrame, text_column: str, tex
             else (shap_values[0].data, shap_values[0].values)
         )
     except Exception as e:
-        logger.exception(f"Failed to explain text: {text_document}", e)
+        logger.error("Failed to explain text %s", text_document)
+        logger.exception(e)
         raise Exception("Failed to create text explanation") from e
 
 
@@ -303,9 +306,12 @@ def text_explanation_prediction_wrapper(
 ) -> Callable:
     def text_predict(text_documents: List[str]):
         num_documents = len(text_documents)
-
-        df_with_text_documents = input_example.append([input_example] * (num_documents - 1), ignore_index=True)
-        df_with_text_documents[text_column] = pd.DataFrame(text_documents)
+        df_with_text_documents = (
+            input_example.copy()
+            if num_documents == 1
+            else pd.concat([input_example] * num_documents, ignore_index=True)
+        )
+        df_with_text_documents[text_column] = text_documents
         return prediction_function(df_with_text_documents)
 
     return text_predict

giskard/scanner/llm/gender_stereotype_detector.py

@@ -1,6 +1,7 @@
+from typing import Sequence
+
 import re
 from pathlib import Path
-from typing import Sequence
 
 import pandas as pd
 import scipy.stats as stats
@@ -21,12 +22,11 @@ def __init__(self, threshold: float = 0.05):
     def run(self, model: LangchainModel, dataset: Dataset) -> Sequence[Issue]:
         # @TODO: add Winogender Schemas
         df_job = self._read_job_data()
-
-        dataset = Dataset(df=df_job.loc[:, ("job",)], column_types={"job": "text"})
+        read_dataset = Dataset(df=df_job.loc[:, ("job",)], column_types={"job": "text"})
         test_model = model.rewrite_prompt(_prompt_template, input_variables=["job"])
 
         # Get model output and count gender-specific pronouns
-        output = test_model.predict(dataset).prediction
+        output = test_model.predict(read_dataset).prediction
         detected_genders = [detect_gender(sentence) for sentence in output]
 
         df = df_job.copy()
@@ -64,7 +64,7 @@ def run(self, model: LangchainModel, dataset: Dataset) -> Sequence[Issue]:
             return [
                 Issue(
                     model,
-                    dataset,
+                    read_dataset,
                     level=IssueLevel.MAJOR,
                     group=Stereotypes,
                     description=desc,

giskard/scanner/llm/harmfulness_detector.py

@@ -1,6 +1,7 @@
-from pathlib import Path
 from typing import List, Sequence
 
+from pathlib import Path
+
 import pandas as pd
 
 from ...datasets import Dataset
@@ -27,11 +28,11 @@ def run(self, model: LangchainModel, dataset: Dataset) -> Sequence[Issue]:
         )
 
         # Prepare test model and dataset
-        dataset = Dataset(df=prompts_df, column_types={"text": "text"})
+        read_dataset = Dataset(df=prompts_df, column_types={"text": "text"})
         test_model = model.rewrite_prompt("{text}", input_variables=["text"])
 
         # Run prediction and evaluate toxicity/harmfulness
-        output = test_model.predict(dataset).prediction
+        output = test_model.predict(read_dataset).prediction
         harmfulness = self._compute_harmfulness(output)
 
         # Filter the examples based
@@ -58,7 +59,7 @@ def run(self, model: LangchainModel, dataset: Dataset) -> Sequence[Issue]:
         return [
             Issue(
                 model,
-                dataset,
+                read_dataset,
                 level=IssueLevel.MAJOR,
                 group=Harmfulness,
                 description=desc,

giskard/scanner/llm/minority_stereotype_detector.py

@@ -1,6 +1,7 @@
+from typing import List, Optional, Sequence
+
 import itertools
 from pathlib import Path
-from typing import List, Optional, Sequence
 
 import pandas as pd
 
@@ -34,12 +35,12 @@ def run(self, model: LangchainModel, dataset: Dataset) -> Sequence[Issue]:
             ],
             columns=["text", "target"],
         )
-        dataset = Dataset(prompt_df.loc[:, ("text",)], column_types={"text": "text"})
+        read_dataset = Dataset(prompt_df.loc[:, ("text",)], column_types={"text": "text"})
 
         test_model = model.rewrite_prompt("{text}", input_variables=["text"])
 
         # Generate output and predict score
-        output = test_model.predict(dataset).prediction
+        output = test_model.predict(read_dataset).prediction
         bias_score = self._compute_bias(output)
 
         examples = pd.DataFrame(
@@ -63,7 +64,7 @@ def run(self, model: LangchainModel, dataset: Dataset) -> Sequence[Issue]:
             issues.append(
                 Issue(
                     model,
-                    dataset,
+                    read_dataset,
                     level=IssueLevel.MAJOR,
                     group=Stereotypes,
                     meta={
@@ -84,4 +85,4 @@ def _compute_bias(self, sentences: List[str]):
             raise LLMImportError() from err
 
         results = Detoxify().predict(list(sentences))
-        return results["identity_attack"]
+        return results["identity_attack"]