From 8ca103e7605dd79761b66ed8a75d81f126697bbb Mon Sep 17 00:00:00 2001
From: wu <wucaiping1@huawei.com>
Date: Thu, 7 Dec 2023 10:53:59 +0000
Subject: [PATCH] add moderation.py.

Signed-off-by: wu <wucaiping1@huawei.com>
---
 moderation.py | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)
 create mode 100644 moderation.py

diff --git a/moderation.py b/moderation.py
new file mode 100644
index 00000000..661e5d35
--- /dev/null
+++ b/moderation.py
@@ -0,0 +1,55 @@
+import json
+import logging
+import requests
+import redis
+
+from pilot.configs.model_config import config_parser
+
+logger = logging.getLogger("model_worker")
+
+
+class Moderation:
+    def __init__(self) -> None:
+        self.configs = config_parser
+        self.token_url = self.configs.get('moderation', 'token_url')
+        self.moderation_url = self.configs.get('moderation', 'moderation_url')
+        self.redis_host = self.configs.get('redis', 'redis_host')
+        self.redis_port = self.configs.getint('redis', 'redis_port')
+        self.redis_db = self.configs.getint('redis', 'redis_db')
+        self.redis_password = self.configs.get('redis', 'redis_password')
+        self.redis = redis.Redis(host=self.redis_host, port=self.redis_host, db=self.redis_db, password=self.redis_password)
+
+    def get_token(self, token_data):
+        token = self.redis.get('moderation_token')
+        if token:
+            return token
+
+        headers = {
+            "Content-Type": "application/json;charset=utf-8",
+        }
+             
+        response = requests.post(url=self.token_url, data=json.dumps(token_data), headers=headers)
+        token = response.headers.get("X-Subject-Token")
+        created_at = '2023-13-01'
+        self.redis.set('moderation_token', token, 36000)
+        logger.info('...refresh token succeed...')
+        return token
+
+    def check_text(self, text):
+        token = self.get_token(self.username, self.password, self.domain, self.project_id)
+        body = {
+            "categories": ["terrorism"],
+            "items": [{
+                "text": text,
+                "type": "content"
+            }
+            ]
+        }
+        _header = {
+            "Content-Type": 'application/json',
+            'X-Auth-Token': token
+        }
+        reslut = requests.post(url=self.moderation_url, data=json.dumps(body), headers=_header)
+        if reslut.json().get('result').get('suggestion') == 'pass':
+            return True      
+        return False
-- 
Gitee