diff --git a/moderation.py b/moderation.py new file mode 100644 index 0000000000000000000000000000000000000000..05c599d723b04e0e708bc4f1b982da5be72b587b --- /dev/null +++ b/moderation.py @@ -0,0 +1,59 @@ +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.load_config(self.configs) + 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-12-01' + self.redis.set('moderation_token', token, 36000) + logger.info('...refresh token succeed...') + return token + + def check_text(self, text): + token = self.get_token(token_data) + body = { + "categories": ["terrorism"], + "items": [{ + "text": text, + "type": "content" + } + ] + } + _header = { + "Content-Type": 'application/json', + 'X-Auth-Token': token + } + result = requests.get(url=self.moderation_url, data=json.dumps(body), headers=_header) + if result.json().get('result').get('suggestion') == 'pass': + return True + return False + + def load_config(configs): + self.token_url = configs.get('moderation', 'token_url') + self.moderation_url = configs.get('moderation', 'moderation_url') + self.redis_host = configs.get('redis', 'redis_host') + self.redis_port = configs.getint('redis', 'redis_port') + self.redis_db = configs.get('redis', 'redis_db') + self.redis_password = configs.get('redis', 'redis_password') + return self.token_url, self.moderation_url, self.redis_host, self.redis_port, self.redis_db, self.redis_password \ No newline at end of file