Use local DB instead of the Cloud DB

dialogues/bitod/main.py

@@ -1,36 +1,68 @@
 import logging
+import os
 import re
 
 from pymongo import MongoClient
 
 from ..main import WOZDataset
 from .src.knowledgebase import api
 from .src.knowledgebase.en_zh_mappings import BitodMapping
+from ..utils import read_jsonl_files_in_folder
 
 logger = logging.getLogger(__name__)
 
 
+def build_bitod_db(db_json_path, api_map, mongodb_host=""):
+    if mongodb_host:
+        db_client = MongoClient(mongodb_host)
+    else:
+        db_client = MongoClient()
+    bitod_db = db_client["bilingual_tod"]
+    for db in bitod_db.list_collection_names():
+        bitod_db[db].drop()
+
+    raw_db = read_jsonl_files_in_folder(db_json_path)
+    for domain in raw_db.keys():
+        if api_map is None:
+            col = bitod_db[domain]
+        else:
+            col = bitod_db[api_map[domain]]
+        for i in range(len(raw_db[domain])):
+            slot_list = list(raw_db[domain][i].keys())
+            for s in slot_list:
+                if "." in s:
+                    # escape dot cause mongodb doesn't like '.' and '$' in key names
+                    raw_db[domain][i][s.replace(".", "\uFF0E")] = raw_db[domain][i].pop(s)
+        col.insert_many(raw_db[domain], ordered=True)
+    return bitod_db
+
+
+
 class Bitod(WOZDataset):
     def __init__(self, name='bitod'):
         super().__init__(name)
 
         self.value_mapping = BitodMapping()
-
-        mongodb_host = 'mongodb+srv://bitod:plGYPp44hASzGbmm@cluster0.vo7pq.mongodb.net/bilingual_tod?retryWrites=true&w=majority&ssl=true&ssl_cert_reqs=CERT_NONE'
-        client = MongoClient(mongodb_host, authSource='admin')
-        database = client["bilingual_tod"]
+
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        mongodb_host = "mongodb://localhost:27017/"
+        bitod_db = build_bitod_db(
+            db_json_path=os.path.join(*[cur_dir, f'database/db/bilingual_tod']),
+            api_map=None,
+            mongodb_host=mongodb_host,
+        )
 
         db = {"null": None}
         self.skipped_entities = set()
 
         for domain in ['restaurants', 'hotels']:
             for lang in ['en_US', 'zh_CN']:
-                db[f"{domain}_{lang}_booking"] = database[f"{domain}_{lang}"]
-                db[f"{domain}_{lang}_search"] = database[f"{domain}_{lang}"]
+                db[f"{domain}_{lang}_booking"] = bitod_db[f"{domain}_{lang}"]
+                db[f"{domain}_{lang}_search"] = bitod_db[f"{domain}_{lang}"]
 
         for domain in ['attractions', 'weathers']:
             for lang in ['en_US', 'zh_CN']:
-                db[f"{domain}_{lang}_search"] = database[f"{domain}_{lang}"]
+                db[f"{domain}_{lang}_search"] = bitod_db[f"{domain}_{lang}"]
 
         self.db = db
 

dialogues/risawoz/main.py

@@ -1,22 +1,55 @@
 import logging
+import os
 
 from pymongo import MongoClient
 
 from ..main import WOZDataset
 from .src.knowledgebase import api
 from .src.knowledgebase.en_zh_mappings import RisawozMapping
+from ..utils import read_json_files_in_folder
 
 logger = logging.getLogger(__name__)
 
 
+def build_risawoz_db(db_json_path, api_map, mongodb_host=""):
+    if mongodb_host:
+        db_client = MongoClient(mongodb_host)
+    else:
+        db_client = MongoClient()
+    risawoz_db = db_client["risawoz"]
+    for db in risawoz_db.list_collection_names():
+        risawoz_db[db].drop()
+
+    for lang in ['en', 'fr', 'hi', 'ko', 'zh', 'enhi']:
+        folder = f'db_{lang}'
+        raw_db = read_json_files_in_folder(os.path.join(db_json_path, folder))
+        for domain in raw_db.keys():
+            if api_map is None:
+                col = risawoz_db[domain]
+            else:
+                col = risawoz_db[api_map[domain]]
+            for i in range(len(raw_db[domain])):
+                slot_list = list(raw_db[domain][i].keys())
+                for s in slot_list:
+                    if "." in s:
+                        # escape dot cause mongodb doesn't like '.' and '$' in key names
+                        raw_db[domain][i][s.replace(".", "\uFF0E")] = raw_db[domain][i].pop(s)
+            col.insert_many(raw_db[domain], ordered=True)
+    return risawoz_db
+
+
 class Risawoz(WOZDataset):
     def __init__(self, name='risawoz', src='zh', tgt='en', mongodb_host=None):
         super().__init__(name)
-        if mongodb_host is None:
-            mongodb_host = 'mongodb+srv://bitod:plGYPp44hASzGbmm@cluster0.vo7pq.mongodb.net/risawoz?retryWrites=true&w=majority&ssl=true&ssl_cert_reqs=CERT_NONE'
-        client = MongoClient(mongodb_host, authSource='admin')
 
-        self.db = client["risawoz"]
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        mongodb_host = "mongodb://localhost:27017/"
+        risawoz_db = build_risawoz_db(
+            db_json_path=os.path.join(*[cur_dir, f'database/']),
+            api_map=None,
+            mongodb_host=mongodb_host,
+        )
+        self.db = risawoz_db
 
         tgt = tgt.split('_', 1)[0]
         src = src.split('_', 1)[0]

dialogues/risawoz/src/convert.py

@@ -3,28 +3,17 @@
 import json
 import os
 from collections import defaultdict
-from contextlib import ExitStack
-from pathlib import Path
 
 import pymongo
 import requests
 from tqdm.autonotebook import tqdm
 
 from dialogues.risawoz.main import Risawoz
 from dialogues.risawoz.src.knowledgebase.api import call_api, process_string
+from dialogues.utils import read_json_files_in_folder
 
 
-def read_json_files_in_folder(path):
-    json_filename = [path + '/' + filename for filename in os.listdir(path) if '.json' in filename]
-    with ExitStack() as stack:
-        files = [stack.enter_context(open(fname)) for fname in json_filename]
-        data = {}
-        for i in range(len(files)):
-            data[Path(json_filename[i]).stem] = json.load(files[i])
-    return data
-
-
-def build_db(db_json_path, api_map, setting, value_mapping, mongodb_host=""):
+def build_db(db_json_path, api_map, setting, mongodb_host=""):
     raw_db = read_json_files_in_folder(db_json_path)
     if mongodb_host:
         db_client = pymongo.MongoClient(mongodb_host)
@@ -327,10 +316,8 @@ def build_dataset(original_data_path, db, setting, value_mapping, debug=False, m
         db_json_path=os.path.join(*[args.root, f'database/db_{args.setting}']),
         api_map=None,
         setting=args.setting,
-        value_mapping=dataset.value_mapping,
         mongodb_host=mongodb_host,
     )
-    db_client = pymongo.MongoClient(mongodb_host)
 
     # download original RiSAWOZ dataset
     original_data_path = os.path.join(*[args.root, args.data_dir])

dialogues/utils.py

@@ -1,6 +1,9 @@
+import json
 import os
 import re
 import subprocess
+from contextlib import ExitStack
+from pathlib import Path
 
 from word2number import w2n
 
@@ -179,3 +182,26 @@ def constraint_list_to_dict(constraints):
             else:
                 result[name] = constraint_and(result[name], constraint_function)
     return result
+
+
+def read_json_files_in_folder(path):
+    json_filename = [path + '/' + filename for filename in os.listdir(path) if '.json' in filename]
+    with ExitStack() as stack:
+        files = [stack.enter_context(open(fname)) for fname in json_filename]
+        data = {}
+        for i in range(len(files)):
+            data[Path(json_filename[i]).stem] = json.load(files[i])
+    return data
+
+def read_jsonl_files_in_folder(path):
+    json_filename = [path + '/' + filename for filename in os.listdir(path) if '.jsonl' in filename]
+    with ExitStack() as stack:
+        files = [stack.enter_context(open(fname)) for fname in json_filename]
+        data = {}
+        for i in range(len(files)):
+            json_list = list(files[i])
+            result = []
+            for json_str in json_list:
+                result.append(json.loads(json_str))
+            data[Path(json_filename[i]).stem] = result
+    return data

tests/risawoz/test_convert.py

@@ -31,7 +31,6 @@
     db_json_path=os.path.join(*[args.root, f'database/db_{args.setting}']),
     api_map=None,
     setting=args.setting,
-    value_mapping=dataset.value_mapping,
     mongodb_host=mongodb_host,
 )