add model-config.json

main.py

@@ -1,5 +1,6 @@
 import os
 import requests
+import json
 
 from model.bow import BowModel
 
@@ -8,83 +9,77 @@
 BASE_API_URL = 'https://ai-api-stag.cofacts.org/v1'
 TEST = True
 
-DEFAULT_CATEGORY_MAPPING = {
-    0: 'kj287XEBrIRcahlYvQoS',  # 中國影響力
-    1: 'kz3c7XEBrIRcahlYxAp6',  # 性少數與愛滋病
-    2: 'lD3h7XEBrIRcahlYeQqS',  # 女權與性別刻板印象
-    3: 'lT3h7XEBrIRcahlYugqq',  # 保健秘訣、食品安全
-    4: 'lj2m7nEBrIRcahlY6Ao_',  # 基本人權問題
-    5: 'lz2n7nEBrIRcahlYDgri',  # 農林漁牧政策
-    6: 'mD2n7nEBrIRcahlYLAr7',  # 能源轉型
-    7: 'mT2n7nEBrIRcahlYTArI',  # 環境生態保護
-    8: 'mj2n7nEBrIRcahlYdArf',  # 優惠措施、新法規、政策宣導
-    9: 'mz2n7nEBrIRcahlYnQpz',  # 科技、資安、隱私
-    10: 'nD2n7nEBrIRcahlYwQoW',  # 免費訊息詐騙
-    11: 'nT2n7nEBrIRcahlY6QqF',  # 有意義但不包含在以上標籤
-    12: 'nj2n7nEBrIRcahlY-gpc',  # 無意義
-    13: 'nz2o7nEBrIRcahlYBgqQ',  # 廣告
-    14: 'oD2o7nEBrIRcahlYFgpm',  # 只有網址其他資訊不足
-    15: 'oT2o7nEBrIRcahlYKQoM',  # 政治、政黨
-    16: 'oj2o7nEBrIRcahlYRAox'  # 轉發協尋、捐款捐贈
-}
+config = json.loads(open('./model-config.json', 'r').read())
+
+DEFAULT_CATEGORY_MAPPING = dict([(int(k), v)
+                                 for k, v in config['categoryMapping'].items()])
 
 
 def predict():
-  model = BowModel()
-  while True:
-    get_tasks_url = f'{BASE_API_URL}/tasks?modelId={MODEL_ID}&apiKey={API_KEY}'
-    if TEST: get_tasks_url += '&test=1'
+    model = BowModel()
+    while True:
+        get_tasks_url = f'{BASE_API_URL}/tasks?modelId={MODEL_ID}&apiKey={API_KEY}'
+        if TEST:
+            get_tasks_url += '&test=1'
+
+        tasks = requests.get(get_tasks_url).json()
+
+        # print(tasks)
 
-    tasks = requests.get(get_tasks_url).json()
+        if len(tasks) == 0:
+            break
 
-    # print(tasks)
+        result = []
 
-    if len(tasks) == 0: break
+        count = 0
 
-    result = []
+        for task in tasks:
+            text = task['content']
+            count += 1
 
-    count = 0
+            # to avoid oversized requests
+            if count > 100:
+                break
 
-    for task in tasks:
-      text = task['content']
-      count += 1
+            category = DEFAULT_CATEGORY_MAPPING[model.predict_text(text)[0]]
 
-      # to avoid oversized requests
-      if count > 100: break
+            temp = {
+                'id': task['id'],
+                'result': {
+                    'prediction': {
+                        'confidence': {}
+                    },
+                    'time': 1500000
+                }
+            }
+            temp['result']['prediction']['confidence'][category] = 1.0
+            result.append(temp)
 
-      category = DEFAULT_CATEGORY_MAPPING[model.predict_text(text)[0]]
+        send_result = requests.post(f'{BASE_API_URL}/tasks', json=result)
+        print(send_result.text)
 
-      temp = {
-        'id': task['id'],
-        'result': {
-          'prediction': {
-            'confidence': {}
-          },
-          'time': 1500000
-        }
-      }
-      temp['result']['prediction']['confidence'][category] = 1.0
-      result.append(temp)
+        if TEST:
+            break
 
-    send_result = requests.post( f'{BASE_API_URL}/tasks', json=result)
-    print(send_result.text)
-
-    if TEST: break
 
 def register():
-  register_url = f'{BASE_API_URL}/models'
+    register_url = f'{BASE_API_URL}/models'
 
-  result = requests.post(register_url, data={
-    "name": "rumors-ai-bow",
-    "realTime": False,
-    "categoryMapping": DEFAULT_CATEGORY_MAPPING
-  })
+    result = requests.post(register_url, data={
+        "name": "rumors-ai-bow",
+        "realTime": False,
+        "categoryMapping": DEFAULT_CATEGORY_MAPPING
+    })
+
+    print(result)
 
-  print(result)
 
 def main():
-  if os.getenv('CFA_ACTION') == 'register': register()
-  else: predict()
-
+    if os.getenv('CFA_ACTION') == 'register':
+        register()
+    else:
+        predict()
+
+
 if __name__ == '__main__':
-  main()
+    main()

model-config.json

@@ -0,0 +1,21 @@
+{
+  "categoryMapping": {
+    "0": "kj287XEBrIRcahlYvQoS",
+    "1": "kz3c7XEBrIRcahlYxAp6",
+    "2": "lD3h7XEBrIRcahlYeQqS",
+    "3": "lT3h7XEBrIRcahlYugqq",
+    "4": "lj2m7nEBrIRcahlY6Ao_",
+    "5": "lz2n7nEBrIRcahlYDgri",
+    "6": "mD2n7nEBrIRcahlYLAr7",
+    "7": "mT2n7nEBrIRcahlYTArI",
+    "8": "mj2n7nEBrIRcahlYdArf",
+    "9": "mz2n7nEBrIRcahlYnQpz",
+    "10": "nD2n7nEBrIRcahlYwQoW",
+    "11": "nT2n7nEBrIRcahlY6QqF",
+    "12": "nj2n7nEBrIRcahlY-gpc",
+    "13": "nz2o7nEBrIRcahlYBgqQ",
+    "14": "oD2o7nEBrIRcahlYFgpm",
+    "15": "oT2o7nEBrIRcahlYKQoM",
+    "16": "oj2o7nEBrIRcahlYRAox" 
+  }
+}