Skip to content
GitLab
Switch branch/tag
Find file BlameHistoryPermalink
plugin.py 5.76 KiB
1 
import json
2 
import time
3 
from typing import Any, Dict, List
4
5 
from loguru import logger
6 
from pymemri.data.schema import CategoricalPrediction, Edge, EmailMessage, Item, Message
7 
from pymemri.data.schema import Model as ModelItem
8 
from pymemri.data.schema import Trigger, Tweet
9 
from pymemri.plugin.trigger_plugin_base import TriggerPluginBase
10 
from pymemri.webserver.models.trigger import TriggerReq
11
12 
from .model import Model
13 
from .utils import batch_list, item_to_data
14
15 
TRIGGER_BATCH_SIZE = 16
16
17 
class ClassifierPlugin(TriggerPluginBase):
18
19 
    schema_classes = [Message, CategoricalPrediction, ModelItem, Tweet, Trigger]
20
21 
    def __init__(
22 
        self,
23 
        item_type: str = "Message",
24 
        item_service: str = None,
25 
        model_name: str = "test_release_3_model",
26 
        model_version: str = "0.1",
27 
        isMock: bool = True,
28 
        **kwargs,
29 
    ):
30 
        """
31 
        ClassifierPlugin is a plugin that wraps any classifier and handles all communication with the Pod and conversion from/to `Item`s
32
33 
        Args:
34 
            item_type (str, optional): The Item type this plugin should make predictions on. Defaults to "Message".
35 
            item_service (str, optional): The service of Items this plugin should make predictions on. Defaults to None.
36 
            model_name (str, optional): Name of the model the plugin should use. Defaults to None.
37 
            model_version (str, optional): Version of the model the plugin should use. Defaults to "0.1".
38 
        """
39 
        super().__init__(**kwargs)
40 
        self.model_name = model_name
41 
        self.model_version = model_version
42 
        self.item_type = item_type
43
44
45 
        self.query = {"type": item_type}
46 
        if item_service is not None:
47 
            self.query["service"] = item_service
48 
        if isMock is True:
49 
            self.query["isMock"] = True
50
51
52
53
54 
    def run(self):
55 
        logger.info("Loading model...")
56 
        self.load_model()
57 
        self.set_run_status("daemon")
58 
        logger.info("Done setting up.")
59
60
61 
    def predict(self, item_batch: List[Item]):
62 
        item_batch = self.filter_items(item_batch)
63 
        prepared_batch = self.prepare_batch(item_batch)
64 
        predictions = self.model.predict(prepared_batch)
65 
        prediction_items = [self.prediction_to_item(p) for p in predictions]
66 
        self.sync_to_pod(item_batch, prediction_items)
67
68 
    def load_model(self):
69 
        self.model = Model(client=self.client)
70
7172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
# Without model name, do not create a model item if self.model_name is None: return # Search in pod for existing models with same name and version, # add a new model if it does not exist. model_items = self.client.search( {"type": "Model", "name": self.model_name, "version": self.model_version} ) if model_items: self.model_item = model_items[0] else: self.model_item = ModelItem( name=self.model_name, version=self.model_version ) self.client.create(self.model_item) def filter_items(self, items: List[Item]) -> List[Item]: result = [] for item in items: if item_to_data(item, self.client) is not None: result.append(item) return result def prepare_batch(self, batch: List[Item]) -> List[Any]: """Prepare a list of items for the model. See `utils.item_to_data` for more information. Args: batch (List[Item]): List of Items from the Pod. Returns: List[Any]: List of prepared data. """ return [item_to_data(item, self.client) for item in batch] def prediction_to_item( self, prediction: List[Dict[str, Any]] ) -> CategoricalPrediction: """Converts a prediction returned by self.model to a CategoricalPrediction that can be added to the Pod Args: prediction (List[Dict[str, float]]): List of predictions. For the correct format, see `.model.Model`. Returns: CategoricalPrediction: List of formatted `CategoricalPrediction`s """ # Get the label with the highest score max_label = max(prediction, key=lambda p: p["score"])["label"] return CategoricalPrediction( probs=json.dumps(prediction), value=max_label, source=f"{self.model.name}:{self.model.version}", ) def sync_to_pod(self, items: List[Item], predictions: List[CategoricalPrediction]): """For each item, prediction, add the prediction to the pod and create an edge called 'label' between Item and Prediction. Args: items (List[Item]): [description] predictions (List[Prediction]): [description] """ # Create edges between item and predictions edges = [ Edge(item, prediction, "label") for item, prediction in zip(items, predictions) ] # Create edges between predictions and model if self.model_item is not None: edges += [
141142143144145146147148149150151152153154155156157158159160161162163164165166167
Edge(prediction, self.model_item, "model") for prediction in predictions ] self.client.bulk_action(create_items=predictions, create_edges=edges) def _model_has_loaded(self) -> bool: if hasattr(self, 'model'): return True return False def _wait_for_model_has_loaded(self) -> None: while not self._model_has_loaded(): time.sleep(1) def trigger(self, req: TriggerReq): if len(req.item_ids) == 0: return items = self.client.search({"ids": req.item_ids}) self._wait_for_model_has_loaded() if len(items) < len(req.item_ids): logger.warning("Could not find all `item_ids` in Pod from received trigger.") for item_batch in batch_list(items, TRIGGER_BATCH_SIZE): self.predict(item_batch)