Repository: wayfair-incubator/extra-model (language: python)

Issue:
Add an interface to use `extra-model` inside of Python
Right now all of the interfaces we expose assume the workflow like this: `csv` file in -> `csv` file out.
It would be great to expose another interface that takes `pandas` df and returns `pandas` df as an output. We do this internally anyways, so nothing super difficult.


Discussion / hints:
Hi can anyone review this PR for this issue https://github.com/wayfair-incubator/extra-model/pull/198, thanks

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2c08920e..3d5155ec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.3.1]
+
+### Changed
+- Added `run_from_dataframe` in `_run.py`, which enables passing dataframe as an input
+
+
 ## [0.3.0]
 
 ### Changed
diff --git a/docs/examples/code/extra_model_inside_python.ipynb b/docs/examples/code/extra_model_inside_python.ipynb
new file mode 100644
index 00000000..8ac78736
--- /dev/null
+++ b/docs/examples/code/extra_model_inside_python.ipynb
@@ -0,0 +1,452 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b90e2f55",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>Comments</th>\n",
+       "      <th>CommentId</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>My salad had a bland vinegrette on the baby gr...</td>\n",
+       "      <td>561</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Appetite instantly gone.</td>\n",
+       "      <td>997</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>The only downside is the service.</td>\n",
+       "      <td>321</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>We waited for forty five minutes in vain.</td>\n",
+       "      <td>153</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>It took over 30 min to get their milkshake, wh...</td>\n",
+       "      <td>74</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>95</th>\n",
+       "      <td>You get incredibly fresh fish, prepared with c...</td>\n",
+       "      <td>296</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>96</th>\n",
+       "      <td>I give it 2 thumbs down</td>\n",
+       "      <td>149</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>97</th>\n",
+       "      <td>AN HOUR... seriously?</td>\n",
+       "      <td>569</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>98</th>\n",
+       "      <td>On the up side, their cafe serves really good ...</td>\n",
+       "      <td>100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>99</th>\n",
+       "      <td>Del Taco is pretty nasty and should be avoided...</td>\n",
+       "      <td>965</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>100 rows × 2 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                             Comments  CommentId\n",
+       "0   My salad had a bland vinegrette on the baby gr...        561\n",
+       "1                            Appetite instantly gone.        997\n",
+       "2                   The only downside is the service.        321\n",
+       "3           We waited for forty five minutes in vain.        153\n",
+       "4   It took over 30 min to get their milkshake, wh...         74\n",
+       "..                                                ...        ...\n",
+       "95  You get incredibly fresh fish, prepared with c...        296\n",
+       "96                            I give it 2 thumbs down        149\n",
+       "97                              AN HOUR... seriously?        569\n",
+       "98  On the up side, their cafe serves really good ...        100\n",
+       "99  Del Taco is pretty nasty and should be avoided...        965\n",
+       "\n",
+       "[100 rows x 2 columns]"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import pandas as pd\n",
+    "df = pd.read_csv('../../../tests/resources/100_comments.csv')\n",
+    "df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "65113718",
+   "metadata": {},
+   "source": [
+    "### We can pass is_dataframe=True and pass input_path as a dataframe"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "b2c2d795",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/lib/python3.9/site-packages/extra_model-0.3.0-py3.9.egg/extra_model/_disambiguate.py:111: RuntimeWarning: invalid value encountered in true_divide\n",
+      "  contexts = [np.divide(context, np.linalg.norm(context)) for context in contexts]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>Position</th>\n",
+       "      <th>Aspect</th>\n",
+       "      <th>Descriptor</th>\n",
+       "      <th>AspectCount</th>\n",
+       "      <th>WordnetNode</th>\n",
+       "      <th>SentimentCompound</th>\n",
+       "      <th>SentimentBinary</th>\n",
+       "      <th>AdCluster</th>\n",
+       "      <th>CommentId</th>\n",
+       "      <th>Topic</th>\n",
+       "      <th>TopicImportance</th>\n",
+       "      <th>TopicSentimentCompound</th>\n",
+       "      <th>TopicSentimentBinary</th>\n",
+       "      <th>TopicCount</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>9</td>\n",
+       "      <td>downside</td>\n",
+       "      <td>only</td>\n",
+       "      <td>1</td>\n",
+       "      <td>downside.n.01</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>only</td>\n",
+       "      <td>321</td>\n",
+       "      <td>downside.n.01</td>\n",
+       "      <td>0.007051</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>65</td>\n",
+       "      <td>chips</td>\n",
+       "      <td>many</td>\n",
+       "      <td>1</td>\n",
+       "      <td>chip.n.04</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>many</td>\n",
+       "      <td>303</td>\n",
+       "      <td>food.n.01</td>\n",
+       "      <td>0.032798</td>\n",
+       "      <td>-0.073582</td>\n",
+       "      <td>-0.090909</td>\n",
+       "      <td>11</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>77</td>\n",
+       "      <td>treat</td>\n",
+       "      <td>real</td>\n",
+       "      <td>1</td>\n",
+       "      <td>dainty.n.01</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>real</td>\n",
+       "      <td>356</td>\n",
+       "      <td>food.n.01</td>\n",
+       "      <td>0.032798</td>\n",
+       "      <td>-0.073582</td>\n",
+       "      <td>-0.090909</td>\n",
+       "      <td>11</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>33</td>\n",
+       "      <td>food</td>\n",
+       "      <td>mediocre</td>\n",
+       "      <td>4</td>\n",
+       "      <td>food.n.01</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>mediocre</td>\n",
+       "      <td>932</td>\n",
+       "      <td>food.n.01</td>\n",
+       "      <td>0.032798</td>\n",
+       "      <td>-0.073582</td>\n",
+       "      <td>-0.090909</td>\n",
+       "      <td>11</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>6</td>\n",
+       "      <td>food</td>\n",
+       "      <td>Worst</td>\n",
+       "      <td>4</td>\n",
+       "      <td>food.n.01</td>\n",
+       "      <td>-0.6249</td>\n",
+       "      <td>-1.0</td>\n",
+       "      <td>Worst</td>\n",
+       "      <td>860</td>\n",
+       "      <td>food.n.01</td>\n",
+       "      <td>0.032798</td>\n",
+       "      <td>-0.073582</td>\n",
+       "      <td>-0.090909</td>\n",
+       "      <td>11</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>72</th>\n",
+       "      <td>51</td>\n",
+       "      <td>trap</td>\n",
+       "      <td>old</td>\n",
+       "      <td>1</td>\n",
+       "      <td>trap.n.02</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>old</td>\n",
+       "      <td>82</td>\n",
+       "      <td>trap.n.02</td>\n",
+       "      <td>0.014101</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>73</th>\n",
+       "      <td>51</td>\n",
+       "      <td>grease trap</td>\n",
+       "      <td>old</td>\n",
+       "      <td>1</td>\n",
+       "      <td>trap.n.02</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>old</td>\n",
+       "      <td>82</td>\n",
+       "      <td>trap.n.02</td>\n",
+       "      <td>0.014101</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>74</th>\n",
+       "      <td>24</td>\n",
+       "      <td>piece</td>\n",
+       "      <td>long</td>\n",
+       "      <td>1</td>\n",
+       "      <td>piece.n.06</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>long</td>\n",
+       "      <td>881</td>\n",
+       "      <td>piece.n.06</td>\n",
+       "      <td>0.007051</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>75</th>\n",
+       "      <td>71</td>\n",
+       "      <td>bit</td>\n",
+       "      <td>little</td>\n",
+       "      <td>1</td>\n",
+       "      <td>bit.n.05</td>\n",
+       "      <td>0.0000</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>little</td>\n",
+       "      <td>984</td>\n",
+       "      <td>bit.n.05</td>\n",
+       "      <td>0.007051</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>0.000000</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>76</th>\n",
+       "      <td>25</td>\n",
+       "      <td>fish</td>\n",
+       "      <td>fresh</td>\n",
+       "      <td>1</td>\n",
+       "      <td>fish.n.02</td>\n",
+       "      <td>0.3182</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>fresh</td>\n",
+       "      <td>296</td>\n",
+       "      <td>fish.n.02</td>\n",
+       "      <td>0.007051</td>\n",
+       "      <td>0.318200</td>\n",
+       "      <td>1.000000</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>77 rows × 14 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "    Position       Aspect Descriptor  AspectCount    WordnetNode  \\\n",
+       "0          9     downside       only            1  downside.n.01   \n",
+       "1         65        chips       many            1      chip.n.04   \n",
+       "2         77        treat       real            1    dainty.n.01   \n",
+       "3         33         food   mediocre            4      food.n.01   \n",
+       "4          6         food      Worst            4      food.n.01   \n",
+       "..       ...          ...        ...          ...            ...   \n",
+       "72        51         trap        old            1      trap.n.02   \n",
+       "73        51  grease trap        old            1      trap.n.02   \n",
+       "74        24        piece       long            1     piece.n.06   \n",
+       "75        71          bit     little            1       bit.n.05   \n",
+       "76        25         fish      fresh            1      fish.n.02   \n",
+       "\n",
+       "    SentimentCompound  SentimentBinary AdCluster  CommentId          Topic  \\\n",
+       "0              0.0000              0.0      only        321  downside.n.01   \n",
+       "1              0.0000              0.0      many        303      food.n.01   \n",
+       "2              0.0000              0.0      real        356      food.n.01   \n",
+       "3              0.0000              0.0  mediocre        932      food.n.01   \n",
+       "4             -0.6249             -1.0     Worst        860      food.n.01   \n",
+       "..                ...              ...       ...        ...            ...   \n",
+       "72             0.0000              0.0       old         82      trap.n.02   \n",
+       "73             0.0000              0.0       old         82      trap.n.02   \n",
+       "74             0.0000              0.0      long        881     piece.n.06   \n",
+       "75             0.0000              0.0    little        984       bit.n.05   \n",
+       "76             0.3182              1.0     fresh        296      fish.n.02   \n",
+       "\n",
+       "    TopicImportance  TopicSentimentCompound  TopicSentimentBinary  TopicCount  \n",
+       "0          0.007051                0.000000              0.000000           1  \n",
+       "1          0.032798               -0.073582             -0.090909          11  \n",
+       "2          0.032798               -0.073582             -0.090909          11  \n",
+       "3          0.032798               -0.073582             -0.090909          11  \n",
+       "4          0.032798               -0.073582             -0.090909          11  \n",
+       "..              ...                     ...                   ...         ...  \n",
+       "72         0.014101                0.000000              0.000000           2  \n",
+       "73         0.014101                0.000000              0.000000           2  \n",
+       "74         0.007051                0.000000              0.000000           1  \n",
+       "75         0.007051                0.000000              0.000000           1  \n",
+       "76         0.007051                0.318200              1.000000           1  \n",
+       "\n",
+       "[77 rows x 14 columns]"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from extra_model._run import run_from_dataframe\n",
+    "from pathlib import Path\n",
+    "\n",
+    "run_from_dataframe(df, embeddings_path='../../../embeddings')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/docs/running_extra.md b/docs/running_extra.md
index 94b3d478..459f98d4 100644
--- a/docs/running_extra.md
+++ b/docs/running_extra.md
@@ -203,3 +203,16 @@ run(
     output_filename=Path("output_filename.csv")
 )
 ```
+
+You can use `run_from_dataframe` function from `extra_model/_run.py`  to use `extra-model` with dataframe 
+
+```python
+from extra_model._run import run_from_dataframe
+from pathlib import Path
+
+df = pd.DataFrame(
+        data=[[1, "test comment"], [2, "test comment 2"]],
+        columns=["CommentId", "Comments"],
+    )
+run_from_dataframe(df)
+```
\ No newline at end of file
diff --git a/extra-model b/extra-model
new file mode 160000
index 00000000..d596d78a
--- /dev/null
+++ b/extra-model
@@ -0,0 +1,1 @@
+Subproject commit d596d78a8f14dec31f51dd335ec15598f3dbb51d
diff --git a/extra_model/_run.py b/extra_model/_run.py
index fe6f4769..e58f6c4b 100644
--- a/extra_model/_run.py
+++ b/extra_model/_run.py
@@ -12,6 +12,35 @@
 logger = logging.getLogger(__name__)
 
 
+def run_from_dataframe(
+    df: pd.core.frame.DataFrame, embeddings_path: Path = MODELS_FOLDER
+) -> pd.core.frame.DataFrame:
+    """
+    Run extra-model with dataframe as an input instead of csv_path.
+
+    :param df: is a dataframe with with 2 columns: CommentId and Comments.
+    :param embeddings_path: path to the embeddings files
+    :return: dataframe of the extramodel results
+    """
+    logging.basicConfig(format="  %(message)s")
+
+    if not {"CommentId", "Comments"}.issubset(df.columns):
+        raise ExtraModelError(
+            f"Input columns must include `CommentId` and `Comments`, \
+        but got {df.columns.to_list()} instead"
+        )
+
+    extra_model = ExtraModel(models_folder=embeddings_path)
+    extra_model.load_from_files()
+
+    logger.info("Running `extra-model`")
+    results_raw = extra_model.predict(comments=df.to_dict("records"))
+    results = pd.DataFrame(results_raw)
+
+    logger.info("Returning results")
+    return results
+
+
 def run(
     input_path: Path,
     output_path: Path,
@@ -21,20 +50,10 @@ def run(
     """Docstring."""
     logging.basicConfig(format="  %(message)s")
 
-    extra_model = ExtraModel(models_folder=embeddings_path)
-    extra_model.load_from_files()
-
     logger.info(f"Loading data from {input_path}")
     input_data = pd.read_csv(input_path)
-    if not {"CommentId", "Comments"}.issubset(input_data.columns):
-        raise ExtraModelError(
-            f"Input columns must include `CommentId` and `Comments`, \
-        but got {input_data.columns.to_list()} instead"
-        )
 
-    logger.info("Running `extra-model`")
-    results_raw = extra_model.predict(comments=input_data.to_dict("records"))
-    results = pd.DataFrame(results_raw)
+    results = run_from_dataframe(input_data, embeddings_path)
 
     if not output_path.exists():
         logger.info(f"Creating folder {output_path}")


Test patch:
diff --git a/tests/test_run.py b/tests/test_run.py
index 6ce2d0b6..28d758fa 100644
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -2,7 +2,7 @@
 import pytest
 
 from extra_model._errors import ExtraModelError
-from extra_model._run import run
+from extra_model._run import run, run_from_dataframe
 
 
 @pytest.fixture
@@ -71,3 +71,25 @@ def test_run__wrong_input_raises_error(mocker, extra_model_mock, pandas_mock):
 
     with pytest.raises(ExtraModelError, match="wrong_column_1"):
         run(mocker.MagicMock(), mocker.MagicMock())
+
+
+def test_run__run_from_dataframe(mocker, extra_model_mock, pandas_mock):
+    df = pd.DataFrame(
+        data=[[1, "test comment"], [2, "test comment 2"]],
+        columns=["CommentId", "Comments"],
+    )
+    run_from_dataframe(df)
+    assert extra_model_mock.called
+    assert pandas_mock.DataFrame.called
+
+
+def test_run_from_dataframe__wrong_input_raises_error(
+    mocker, extra_model_mock, pandas_mock
+):
+    df = pd.DataFrame(
+        data=[[1]],
+        columns=["wrong_column_1"],
+    )
+
+    with pytest.raises(ExtraModelError, match="wrong_column_1"):
+        run_from_dataframe(df)