Merge pull request #107 from LuisMayo/feat/readd-polyglot

Work in progress in readding polyglot support

README.md

@@ -101,6 +101,7 @@ docker run --rm \
 #### Settings
 The following environment variables are honored by objection_engine:
 - oe_bypass_sentiment: If on any value other than the empty string, the sentiment analysis is bypassed
+- oe_sentiment_model: If "polyglot", polyglot will be used for sentiment analisis, if empty or "hf" it would be a huggingface model. When using polyglot models remember to do python -m polyglot download TASK:sentiment2
 - ~~oe_stats_server~~ (not longer supported as of Objection engine 3.3.0): If present, it will be used as the URL to post statistics to. The server responsible should be similar to https://github.com/LuisMayo/simple-server-counter
 - oe_polyglot_models: (docker only) If on polyglot model(s), the data for the model will be downloaded when starting the container.
 - OE_DIRECT_H264_ENCODING: If "true" then it will directly encode the videos in H264 with OpenCV, instead of encoding in mp4v and re-encoding later. This is faster than the default, but it requires running on windows or having a self-compiled OpenCV build in your system. If you don't know what any of this means, don't enable it.

objection_engine/ace_attorney_scene.py

@@ -7,6 +7,7 @@
 from os.path import join
 from os import environ, getenv
 from turtle import pos
+from polyglot.text import Text
 
 from objection_engine.gavel_slam import GavelSlamObject
 from objection_engine.testimony_indicator import TestimonyIndicatorTextObject
@@ -1115,12 +1116,33 @@ def __init__(self, callbacks: dict = None, verify_sprites: bool = False) -> None
         self.has_gone_to_tense_music: bool = False
         self.callbacks = {} if callbacks is None else callbacks
 
+        self._sentiment_analyzer = None
         # Hugging Face sentiment analyzer
-        self._sentiment_analyzer = pipeline(
-            "sentiment-analysis",
-            model=SENTIMENT_MODEL_PATH,
-            tokenizer=SENTIMENT_MODEL_PATH,
-        )
+        if len(getenv("oe_bypass_sentiment", "")) <= 0:
+            model_setting = getenv("oe_sentiment_model", "hf")
+            if model_setting == 'hf':
+                self._sentiment_analyzer = pipeline(
+                    "sentiment-analysis",
+                    model=SENTIMENT_MODEL_PATH,
+                    tokenizer=SENTIMENT_MODEL_PATH,
+                )
+
+    def hf_sentiment(self, text: str):
+        return self._sentiment_analyzer(text)[0]
+
+    def poly_sentiment(self, text: str):
+        poly_text = Text(text)
+        try:
+            polarity = poly_text.polarity
+        except Exception as e:
+            polarity = 0
+        if polarity > 0.35:
+            return {'label': 'positive', 'score': 1.0}
+        # If polarity is -1 there isn't enough information to determine if it's negative therefore we introduce randomness
+        if polarity < -0.35 and (polarity > -1 or random.random() > 0.25):
+            return {'label': 'negative', 'score': 1.0}
+        return {'label': 'positive', 'score': 0.0}
+
 
     def reload_character_data(self, verify_sprites: bool = False):
         self.character_data = load_character_data(verify_sprites=verify_sprites)
@@ -1424,17 +1446,20 @@ def build_from_comments(
                 self.callbacks["on_comment_processed"](i, len(comments), comment)
 
     def get_sentiment(self, text: str):
-        return (
-            [{"label": "neutral", "score": 1.0}]
-            if (len(getenv("oe_bypass_sentiment", "")) > 0)
-            else self._sentiment_analyzer(text)
-        )
+        if len(getenv("oe_bypass_sentiment", "")) > 0:
+            return {"label": "neutral", "score": 1.0}
+        else:
+            model_setting = getenv("oe_sentiment_model", "hf")
+            if model_setting == 'hf':
+                return self.hf_sentiment(text)
+            else:
+                return self.poly_sentiment(text)
 
     def update_pose_for_sentence(
         self, sentence: Sentence, sprites: list[str], manual_score: float = 0.0
     ):
         if manual_score == 0:
-            sentiment: dict = self.get_sentiment(sentence.raw)[0]
+            sentiment: dict = self.get_sentiment(sentence.raw)
         else:
             sentiment: dict = {
                 "label": "positive" if manual_score > 0 else "negative",
@@ -1476,7 +1501,7 @@ def get_boxes_with_pauses(
 
         # Determine if this should have an objection
         if manual_score == 0:
-            text_polarity_data = self.get_sentiment(pg_text.raw)[0]
+            text_polarity_data = self.get_sentiment(pg_text.raw)
             polarity_type = text_polarity_data["label"]
             polarity_confidence = text_polarity_data["score"]
         else:
@@ -1485,7 +1510,7 @@ def get_boxes_with_pauses(
 
         do_objection = (
             (
-                polarity_type == "negative"
+                (polarity_type == "negative" and random() > 0.3)
                 or (polarity_type == "positive" and random() > 0.7)
             )
             and polarity_confidence > 0.5