Skip to content
GitLab
Switch branch/tag
Find file BlameHistoryPermalink
plugin.py 4.29 KiB
1 
import pathlib
2 
from typing import Any, List
3 
import json
4
5 
import numpy as np
6 
from pymemri.data.itembase import Edge, Item
7 
from pymemri.plugin.pluginbase import PluginBase
8 
from pymemri.pod.client import PodClient
9 
from pymemri.data.schema import Message
10 
from transformers import TextClassificationPipeline
11 
from transformers import pipeline as AutoPipeline
12 
import torch
13
14 
from .schema import CategoricalPrediction, CategoricalLabel
15 
import wandb
16
17
18 
def get_predictions(item):
19 
    return [x for x in item.label if isinstance(x, CategoricalPrediction)]
20
21 
class SentimentAnalysis(PluginBase):
22 
    _default_config = {
23 
        "content_field": "content",
24 
        "item_type": "Message",
25 
        "model_name": "cardiffnlp/twitter-xlm-roberta-base-sentiment",
26 
        "model_head": "eelcovdw/memri_sentiment/model_head:latest",
27 
    }
28
29 
    def __init__(self, **kwargs):
30 
        super().__init__(**kwargs)
31 
        if self.pluginRun and self.pluginRun.settings:
32 
            config = json.loads(self.pluginRun.settings)
33 
        else:
34 
            config = dict()
35 
        self.set_config(config)
36
37 
    def set_config(self, config):
38 
        for item, default_value in self._default_config.items():
39 
            value = config.get(item, default_value)
40 
            setattr(self, item, value)
41
42 
    def load_data(self, filter_predicted=True):
43 
        """Load all data from pod of type `self.item_type`.
44
45 
        Args:
46 
            filter_edge (bool, optional): Remove items from result that already have a prediction.
47
48 
        Returns:
49 
            List[Item]: List of messages
50 
        """
51
52 
        data = self.client.search({"type": self.item_type})
53
54 
        if filter_predicted:
55 
            data = [item for item in data if len(get_predictions(item)) == 0]
56
57 
        return data
58
59 
    def load_model_head(self, pipeline: AutoPipeline):
60 
        if self.model_head is None:
61 
            print("Using default model_head")
62 
            return pipeline
63 
        wandb.login(anonymous="must")
64 
        artifact = wandb.Api().artifact(self.model_head)
65 
        mh_path = artifact.download()
66 
        mh_path = next(pathlib.Path(mh_path).glob("*.model"))
67
68 
        try:
69 
            model_head = torch.load(mh_path, map_location=torch.device("cpu"))
70 
            pipeline.model.classifier = model_head
7172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
pipeline.model.eval() except Exception as e: print(f"Could not load custom model head {self.model_head}") print("Using default model_head") return pipeline def infer_sentiment(self, data: List[Any], pipeline: TextClassificationPipeline): text_data = [getattr(item, self.content_field) for item in data] sentiments = pipeline(text_data) for i, sentiment in enumerate(sentiments): sentiment = { "labels": [elem["label"].lower() for elem in sentiment], "probs": [elem["score"] for elem in sentiment], } # Set prediction = argmax(probs) argmax = np.argmax(sentiment["probs"]) sentiment["name"] = sentiment["labels"][argmax] sentiment = {k: (json.dumps(v) if isinstance(v, list) else v) for k, v in sentiment.items()} sentiments[i] = CategoricalPrediction(**sentiment) return sentiments def save_to_pod(self, data: List[Item], sentiments: List[Item]): edges = [ Edge(item, sentiment, "label") for item, sentiment in zip(data, sentiments) ] self.client.bulk_action(create_items=sentiments, create_edges=edges) def run(self): print("Loading data from pod...") data = self.load_data(self.client) if len(data) == 0: print("No data found.") print("Run completed.") return print("Loading model...") pipeline = AutoPipeline( "sentiment-analysis", model=self.model_name, tokenizer=self.model_name, return_all_scores=True, ) pipeline = self.load_model_head(pipeline) print("Inferring sentiment...") sentiments = self.infer_sentiment(data, pipeline) print(f"Adding {len(sentiments)} predictions to pod...") self.save_to_pod(data, sentiments) print("Run completed.") def add_to_schema(self): self.client.add_to_schema( Message, CategoricalLabel, CategoricalPrediction )