Introduce code from Rust monorepo. This is a mess

2026-10-18 07:41:03 +00:00 · 2022-02-28 05:43:54 -05:00
parent 921fb49f9a
commit ba326bfcfd
1058 changed files with 192192 additions and 0 deletions
@@ -0,0 +1,134 @@
+# Copyright (c) 2020, MeetKai Inc.  All rights reserved.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script downloads Text2Sparql data and processes it into NeMo's neural machine translation dataset format.
+
+Text2Sparql data consists of 3 files which are saved to source_data_dir:
+    - train_queries_v3.tsv
+    - test_easy_queries_v3.tsv
+    - test_hard_queries_v3.tsv
+
+After processing, the script saves them to the target_data_dir as:
+    - train.tsv
+    - test_easy.tsv
+    - test_hard.tsv
+
+
+You may run it with:
+
+python import_datasets \
+    --source_data_dir ./text2sparql_src \
+    --target_data_dir ./text2sparql_tgt
+"""
+
+import argparse
+import csv
+import os
+from urllib.request import Request, urlopen
+
+from nemo.collections.nlp.data.data_utils.data_preprocessing import MODE_EXISTS_TMP, if_exist
+from nemo.utils import logging
+
+base_url = "https://m.meetkai.com/public_datasets/knowledge/"
+prefix_map = {
+    "train_queries_v3.tsv": "train.tsv",
+    "test_easy_queries_v3.tsv": "test_easy.tsv",
+    "test_hard_queries_v3.tsv": "test_hard.tsv",
+}
+
+
+def download_text2sparql(infold: str):
+    """Downloads text2sparql train, test_easy, and test_hard data
+
+    Args:
+        infold: save directory path
+    """
+    os.makedirs(infold, exist_ok=True)
+
+    for prefix in prefix_map:
+        url = base_url + prefix
+
+        logging.info(f"Downloading: {url}")
+        if if_exist(infold, [prefix]):
+            logging.info("** Download file already exists, skipping download")
+        else:
+            req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
+            with open(os.path.join(infold, prefix), "wb") as handle:
+                handle.write(urlopen(req, timeout=20).read())
+
+
+def process_text2sparql(infold: str, outfold: str, do_lower_case: bool):
+    """ Process and convert MeetKai's text2sparql datasets to NeMo's neural machine translation format.
+
+    Args:
+        infold: directory path to raw text2sparql data containing
+            train.tsv, test_easy.tsv, test_hard.tsv
+        outfold: output directory path to save formatted data for NeuralMachineTranslationDataset
+            the first line is header (sentence [tab] label)
+            each line should be [sentence][tab][label]
+        do_lower_case: if true, convert all sentences and labels to lower
+    """
+    logging.info(f"Processing Text2Sparql dataset and storing at: {outfold}")
+
+    os.makedirs(outfold, exist_ok=True)
+
+    dataset_name = "Text2Sparql"
+    for prefix in prefix_map:
+        input_file = os.path.join(infold, prefix)
+        output_file = os.path.join(outfold, prefix_map[prefix])
+
+        if if_exist(outfold, [prefix_map[prefix]]):
+            logging.info(f"** {MODE_EXISTS_TMP.format(prefix_map[prefix], dataset_name, output_file)}")
+            continue
+
+        if not if_exist(infold, [prefix]):
+            logging.info(f"** {prefix} of {dataset_name}" f" is skipped as it was not found")
+            continue
+
+        assert input_file != output_file, "input file cannot equal output file"
+        with open(input_file, "r") as in_file:
+            with open(output_file, "w") as out_file:
+                reader = csv.reader(in_file, delimiter="\t")
+
+                # replace headers
+                out_file.write("sentence\tlabel\n")
+                next(reader)
+
+                for line in reader:
+                    sentence = line[0]
+                    label = line[1]
+                    if do_lower_case:
+                        sentence = sentence.lower()
+                        label = label.lower()
+                    out_file.write(f"{sentence}\t{label}\n")
+
+
+if __name__ == "__main__":
+    # Parse the command-line arguments.
+    parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo's format")
+    parser.add_argument(
+        "--source_data_dir", required=True, type=str, help="Path to the folder containing the dataset files"
+    )
+    parser.add_argument("--target_data_dir", required=True, type=str, help="Path to save the processed dataset")
+    parser.add_argument("--do_lower_case", action="store_true")
+    args = parser.parse_args()
+
+    source_dir = args.source_data_dir
+    target_dir = args.target_data_dir
+    do_lower_case = args.do_lower_case
+
+    download_text2sparql(infold=source_dir)
+    process_text2sparql(infold=source_dir, outfold=target_dir, do_lower_case=do_lower_case)