style(black): 🎨 black format notebook

protectai · Apr 6, 2024 · 2f66207 · 2f66207
1 parent 28129d9
commit 2f66207
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 31 deletions.
diff --git a/modelscan/modelscan.py b/modelscan/modelscan.py
@@ -135,7 +135,7 @@ def _iterate_models(self, model_path: Path) -> Generator[Model, None, None]:
                         ModelScanSkipped(
                             "ModelScan",
                             SkipCategories.BAD_ZIP,
-                            "Skipping zip file due to error: %s" % e,
+                            f"Skipping zip file due to error: {e}",
                             str(model.get_source()),
                         )
                     )
@@ -198,8 +198,7 @@ def _scan_source(
                     ModelScanError(
                         scanner.full_name(),
                         ErrorCategories.MODEL_SCAN,
-                        "Error encountered from scanner %s: %s"
-                        % (scanner.full_name(), e),
+                        f"Error encountered from scanner {scanner.full_name()}: {e}",
                         str(model.get_source()),
                     )
                 )
@@ -354,7 +353,7 @@ def generate_report(self) -> Optional[str]:
                 ModelScanError(
                     "ModelScan",
                     ErrorCategories.MODEL_SCAN,
-                    "Error generating report using %s: %s" % (reporting_module, e),
+                    f"Error generating report using {reporting_module}: {e}",
                 )
             )
 

diff --git a/notebooks/pytorch_sentiment_analysis.ipynb b/notebooks/pytorch_sentiment_analysis.ipynb
@@ -16,7 +16,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -56,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -69,10 +69,11 @@
    ],
    "source": [
     "import torch\n",
-    "import os \n",
+    "import os\n",
     "from utils.pytorch_sentiment_model import download_model, predict_sentiment\n",
     "from utils.pickle_codeinjection import PickleInject, get_payload\n",
-    "%env TOKENIZERS_PARALLELISM=false\n"
+    "\n",
+    "%env TOKENIZERS_PARALLELISM=false"
    ]
   },
   {
@@ -87,7 +88,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -112,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -124,7 +125,9 @@
     }
    ],
    "source": [
-    "sentiment = predict_sentiment(\"Stock market was bearish today\", torch.load(safe_model_path))"
+    "sentiment = predict_sentiment(\n",
+    "    \"Stock market was bearish today\", torch.load(safe_model_path)\n",
+    ")"
    ]
   },
   {
@@ -138,16 +141,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "No settings file detected at /Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/modelscan-settings.toml. Using defaults. \n",
+      "No settings file detected at /Users/candiedcode/Code/modelscan/notebooks/modelscan-settings.toml. Using defaults. \n",
       "\n",
-      "Scanning /Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/PyTorchModels/safe_model.pt:safe_model/data.pkl using modelscan.scanners.PickleUnsafeOpScan model scan\n",
+      "Scanning /Users/candiedcode/Code/modelscan/notebooks/PyTorchModels/safe_model.pt:safe_model/data.pkl using modelscan.scanners.PickleUnsafeOpScan model scan\n",
       "\n",
       "\u001b[34m--- Summary ---\u001b[0m\n",
       "\n",
@@ -174,7 +177,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -205,14 +208,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "cat: /Users/jennifercwagenberg/.aws/secrets: No such file or directory\n"
+      "cat: /Users/candiedcode/.aws/secrets: No such file or directory\n"
      ]
     },
     {
@@ -240,16 +243,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "No settings file detected at /Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/modelscan-settings.toml. Using defaults. \n",
+      "No settings file detected at /Users/candiedcode/Code/modelscan/notebooks/modelscan-settings.toml. Using defaults. \n",
       "\n",
-      "Scanning /Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/PyTorchModels/unsafe_model.pt:unsafe_model/data.pkl using modelscan.scanners.PickleUnsafeOpScan model scan\n",
+      "Scanning /Users/candiedcode/Code/modelscan/notebooks/PyTorchModels/unsafe_model.pt:unsafe_model/data.pkl using modelscan.scanners.PickleUnsafeOpScan model scan\n",
       "\n",
       "\u001b[34m--- Summary ---\u001b[0m\n",
       "\n",
@@ -269,7 +272,7 @@
       "Unsafe operator found:\n",
       "  - Severity: CRITICAL\n",
       "  - Description: Use of unsafe operator 'system' from module 'posix'\n",
-      "  - Source: /Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/PyTorchModels/unsafe_model.pt:unsafe_model/data.pkl\n",
+      "  - Source: /Users/candiedcode/Code/modelscan/notebooks/PyTorchModels/unsafe_model.pt:unsafe_model/data.pkl\n",
       "\n",
       "\u001b[34m--- Skipped --- \u001b[0m\n",
       "\n",
@@ -278,7 +281,7 @@
     }
    ],
    "source": [
-    "!modelscan --path  ./PyTorchModels/unsafe_model.pt "
+    "!modelscan --path  ./PyTorchModels/unsafe_model.pt"
    ]
   },
   {
@@ -300,21 +303,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "No settings file detected at /Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/modelscan-settings.toml. Using defaults. \n",
+      "No settings file detected at /Users/candiedcode/Code/modelscan/notebooks/modelscan-settings.toml. Using defaults. \n",
       "\n",
-      "Scanning /Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/PyTorchModels/unsafe_model.pt:unsafe_model/data.pkl using modelscan.scanners.PickleUnsafeOpScan model scan\n",
+      "Scanning /Users/candiedcode/Code/modelscan/notebooks/PyTorchModels/unsafe_model.pt:unsafe_model/data.pkl using modelscan.scanners.PickleUnsafeOpScan model scan\n",
       "\u001b[1m{\u001b[0m\u001b[32m\"summary\"\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m\"total_issues_by_severity\"\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m\"LOW\"\u001b[0m: \u001b[1;36m0\u001b[0m, \u001b[32m\"MEDIUM\"\u001b[0m: \u001b[1;36m0\u001b[0m, \u001b[32m\"HIGH\"\u001b[0m: \u001b[1;36m0\u001b[0m, \n",
       "\u001b[32m\"CRITICAL\"\u001b[0m: \u001b[1;36m1\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m\"total_issues\"\u001b[0m: \u001b[1;36m1\u001b[0m, \u001b[32m\"input_path\"\u001b[0m: \n",
       "\u001b[32m\"./PyTorchModels/unsafe_model.pt\"\u001b[0m, \u001b[32m\"absolute_path\"\u001b[0m: \n",
-      "\u001b[32m\"/Users/jennifercwagenberg/Code/Personal/fork/modelscan/notebooks/PyTorchModels\"\u001b[0m\n",
-      ", \u001b[32m\"modelscan_version\"\u001b[0m: \u001b[32m\"0.0.0\"\u001b[0m, \u001b[32m\"timestamp\"\u001b[0m: \u001b[32m\"2024-04-05T21:31:34.897088\"\u001b[0m, \n",
+      "\u001b[32m\"/Users/candiedcode/Code/modelscan/notebooks/PyTorchModels\"\u001b[0m\n",
+      ", \u001b[32m\"modelscan_version\"\u001b[0m: \u001b[32m\"0.0.0\"\u001b[0m, \u001b[32m\"timestamp\"\u001b[0m: \u001b[32m\"2024-04-06T08:32:44.102497\"\u001b[0m, \n",
       "\u001b[32m\"scanned\"\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m\"total_scanned\"\u001b[0m: \u001b[1;36m1\u001b[0m, \u001b[32m\"scanned_files\"\u001b[0m: \n",
       "\u001b[1m[\u001b[0m\u001b[32m\"unsafe_model.pt:unsafe_model/data.pkl\"\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m\"issues\"\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m\"description\"\u001b[0m: \u001b[32m\"Use of \u001b[0m\n",
       "\u001b[32munsafe operator 'system' from module 'posix'\"\u001b[0m, \u001b[32m\"operator\"\u001b[0m: \u001b[32m\"system\"\u001b[0m, \u001b[32m\"module\"\u001b[0m: \n",

diff --git a/notebooks/utils/pytorch_sentiment_model.py b/notebooks/utils/pytorch_sentiment_model.py
@@ -1,4 +1,4 @@
-from typing import Final
+from typing import Any, Final
 from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 import numpy as np
@@ -42,13 +42,13 @@ def download_model(safe_model_path: str) -> None:
     torch.save(model, safe_model_path)
 
 
-def predict_sentiment(text: str, model: str) -> None:
+def predict_sentiment(text: str, model: Any) -> None:
     """
     Predict the sentiment of a given text using a pre-trained sentiment analysis model.
 
     Args:
         text: The input text to analyze.
-        model: The name or path of the pre-trained sentiment analysis model.
+        model: The sentiment analysis model.
     """
     pretrained_model_name = "cardiffnlp/twitter-roberta-base-sentiment"
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name)