dslighting 1.3.9__py3-none-any.whl → 1.7.6__py3-none-any.whl

dslighting/registry/tensorflow-speech-recognition-challenge/prepare_val.py

@@ -0,0 +1,148 @@
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from tqdm import tqdm
+
+from mlebench.utils import extract
+
+CLASSES = [
+    "yes",
+    "no",
+    "up",
+    "down",
+    "left",
+    "right",
+    "on",
+    "off",
+    "stop",
+    "go",
+    "unknown",
+    "silence",
+]
+
+
+@dataclass(frozen=True)
+class AudioFile:
+    label: str
+    path: Path
+
+
+def _create_and_populate_split(
+    files_to_split: List[AudioFile],
+    public_dir: Path,
+    private_dir: Path,
+    test_size: float,
+    random_state: int,
+) -> List[AudioFile]:
+    """
+    Splits a list of audio files into train and test sets, and populates the
+    corresponding public and private directories.
+
+    Args:
+        files_to_split: A list of AudioFile objects to be split.
+        public_dir: The public directory to save training data and test stubs.
+        private_dir: The private directory to save test ground truth.
+        test_size: The proportion of the dataset to allocate to the test split.
+        random_state: The seed used by the random number generator.
+
+    Returns:
+        A list of AudioFile objects that were assigned to the training set.
+    """
+    train_files, test_files = train_test_split(
+        files_to_split, test_size=test_size, random_state=random_state
+    )
+
+    # Make necessary directories
+    public_dir.mkdir(parents=True, exist_ok=True)
+    private_dir.mkdir(parents=True, exist_ok=True)
+
+    labels = list(
+        dict.fromkeys([file.label for file in train_files])
+    )  # Gets unique elements deterministically
+
+    for label in labels:
+        (public_dir / "train" / "audio" / label).mkdir(parents=True, exist_ok=True)
+
+    (public_dir / "test" / "audio").mkdir(parents=True, exist_ok=True)
+
+    # Copy over train and test files
+    desc_suffix = public_dir.name
+    for file in tqdm(train_files, desc=f"Copying train files to {desc_suffix}"):
+        shutil.copyfile(
+            src=file.path,
+            dst=public_dir / "train" / "audio" / file.label / file.path.name,
+        )
+
+    test_records = []
+
+    for idx, file in enumerate(tqdm(test_files, desc=f"Copying test files to {desc_suffix}")):
+        # Rename files, since training audio files across labels aren't necessarily unique.
+        new_id = str(idx).zfill(8)
+        new_name = f"clip_{new_id}.wav"
+        test_records.append({"fname": new_name, "label": file.label})
+
+        shutil.copyfile(
+            src=file.path,
+            dst=public_dir / "test" / "audio" / new_name,
+        )
+
+    test = pd.DataFrame.from_records(test_records)
+    test.to_csv(private_dir / "test.csv", index=False)
+
+    test_without_labels = test.drop(columns=["label"])
+    sample_submission = test_without_labels.copy()
+    sample_submission["label"] = "silence"
+    sample_submission.to_csv(public_dir / "sample_submission.csv", index=False)
+
+    # Sanity checks
+    test_audio_files = list((public_dir / "test" / "audio").glob("*.wav"))
+    num_test_files = len(test_audio_files)
+    num_submission_entries = len(sample_submission)
+    assert num_test_files == num_submission_entries, (
+        f"The number of test audio files ({num_test_files}) does not match the number of entries "
+        f"in sample_submission.csv ({num_submission_entries}). Please check the file copying process."
+    )
+
+    return train_files
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    # extract only what we need
+    extract(raw / "train.7z", raw)
+
+    # Create train, test from train split
+    audio_dir = raw / "train" / "audio"
+    all_audio_files = sorted(
+        [AudioFile(fpath.parent.name, fpath) for fpath in audio_dir.rglob("*.wav")],
+        key=lambda x: f"{x.label}_{x.path.name}",
+    )
+
+    # --- First Split: Create the original train/test sets ---
+    # This call populates the original `public` and `private` directories.
+    # It returns the files assigned to the training set for the next split.
+    train_files_for_val_split = _create_and_populate_split(
+        files_to_split=all_audio_files,
+        public_dir=public,
+        private_dir=private,
+        test_size=0.1,
+        random_state=0,
+    )
+
+    # --- Second Split: Create the new train_val/test_val sets ---
+    # Define new paths for the validation split, parallel to the original ones.
+    public_val = public.parent / "public_val"
+    private_val = private.parent / "private_val"
+
+    # To get a validation test set of the same size as the original test set,
+    # we take a 1/9 slice of the remaining 90% of the data (0.1 / 0.9 = 1/9).
+    _create_and_populate_split(
+        files_to_split=train_files_for_val_split,
+        public_dir=public_val,
+        private_dir=private_val,
+        test_size=1 / 9,
+        random_state=0, # Use the same random_state for deterministic behavior
+    )

dslighting/registry/tensorflow2-question-answering/checksums.yaml

@@ -0,0 +1,8 @@
+private:
+  gold_submission.csv: 544ba606cfb641eed0d222c3d94ffee6
+  test.jsonl: 86393efdae1c319a8d1d18e36f490e97
+public:
+  sample_submission.csv: 483bdebc6d8e55344c3e0124061eb0ff
+  simplified-nq-test.jsonl: 3abc9e42896e04f4aeaa493e5944e49a
+  simplified-nq-train.jsonl: 47ef1135698eb044006d1e65b68916b5
+zip: 4d3bbcfad1ff7ea12f3c1c3347c4f821

dslighting/registry/tensorflow2-question-answering/config.yaml

@@ -0,0 +1,29 @@
+id: tensorflow2-question-answering
+name: TensorFlow 2.0 Question Answering
+competition_type: code
+awards_medals: true
+prizes:
+  - position: 1
+    value: 12000
+  - position: 2
+    value: 8000
+  - position: 3
+    value: 5000
+  - position: TensorFlow 2.0 Prize 1
+    value: 12000
+  - position: TensorFlow 2.0 Prize 2
+    value: 8000
+  - position: TensorFlow 2.0 Prize 3
+    value: 5000
+description: mlebench/competitions/tensorflow2-question-answering/description.md
+
+dataset:
+  answers: tensorflow2-question-answering/prepared/private/test.jsonl
+  gold_submission: tensorflow2-question-answering/prepared/private/gold_submission.csv
+  sample_submission: tensorflow2-question-answering/prepared/public/sample_submission.csv
+
+grader:
+  name: micro-f1-score
+  grade_fn: mlebench.competitions.tensorflow2-question-answering.grade:grade
+
+preparer: mlebench.competitions.tensorflow2-question-answering.prepare:prepare

dslighting/registry/tensorflow2-question-answering/description.md

@@ -0,0 +1,173 @@
+# Overview
+
+## Description
+
+*"Why is the sky blue?"*
+
+This is a question an [open-domain question answering](https://en.wikipedia.org/wiki/Question_answering#Open_domain_question_answering) (QA) system should be able to respond to. QA systems emulate how people look for information by reading the web to return answers to common questions. Machine learning can be used to improve the accuracy of these answers.
+
+Existing natural language models have been focused on extracting answers from a short paragraph rather than reading an entire page of content for proper context. As a result, the responses can be complicated or lengthy. A good answer will be both succinct and relevant.
+
+In this competition, your goal is to predict short and long answer responses to real questions about Wikipedia articles. The dataset is provided by [Google's Natural Questions](https://ai.google.com/research/NaturalQuestions/dataset), but contains its own unique private test set. A [visualization of examples](https://ai.google.com/research/NaturalQuestions/visualization) shows long and---where available---short answers. In addition to prizes for the top teams, there is a special set of awards for using TensorFlow 2.0 APIs.
+
+If successful, this challenge will help spur the development of more effective and robust QA systems.
+
+## About TensorFlow
+
+TensorFlow is an open source platform for machine learning. With TensorFlow 2.0, tf.keras is the preferred high-level API for TensorFlow, to make model building easier and more intuitive. You may use the [tf.keras built-in compile()/fit() methods](https://www.tensorflow.org/guide/keras/train_and_evaluate#part_i_using_build-in_training_evaluation_loops), or write your own [custom training loops](https://www.tensorflow.org/guide/keras/train_and_evaluate#part_ii_writing_your_own_training_evaluation_loops_from_scratch). See the [Effective TensorFlow 2.0 guide](https://www.tensorflow.org/guide/effective_tf2) and the [tf.keras guide](https://www.tensorflow.org/guide/keras/) for more details.
+
+TensorFlow 2.0 was [recently released](https://medium.com/tensorflow/tensorflow-2-0-is-now-available-57d706c2a9ab) and this competition is to challenge Kagglers to use TensorFlow 2.0's APIs focused on usability, and easier, more intuitive development, to make advancements on Question Answering.
+
+## Evaluation
+
+Submissions are evaluated using [micro F1](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html) between the predicted and expected answers. Predicted long and short answers must match exactly the token indices of one of the ground truth labels ((or match YES/NO if the question has a yes/no short answer). There may be up to five labels for long answers, and more for short. If no answer applies, leave the prediction blank/null.
+
+The metric in this competition diverges from the [original metric](https://github.com/google-research-datasets/natural-questions/blob/master/nq_eval.py) in two key respects: 1) short and long answer formats do not receive separate scores, but are instead combined into a [micro F1 ](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html)score across both formats, and 2) this competition's metric does *not* use `confidence` scores to find an optimal threshold for predictions.
+
+Additional detail:
+
+$F_1=2 * \frac{\text { precision } * \text { recall }}{\text { precision }+ \text { recall }}$
+
+where:
+
+$\begin{gathered}\text { precision }=\frac{T P}{T P+F P} \\ \text { recall }=\frac{T P}{T P+F N}\end{gathered}$
+
+and:
+
+> TP = the predicted indices match one of the possible ground truth indices\
+> FP = the predicted indices do NOT match one of the possible ground truth indices, OR a prediction has been made where no ground truth exists\
+> FN = no prediction has been made where a ground truth exists
+
+In "micro" F1, both long and short answers count toward an overall precision and recall, which is then used to calculate a single F1 value. (In contrast, in "macro" F1 a separate F1 score is calculated for each class / label and then averaged.)
+
+## Submission File
+
+For each ID in the test set, you must predict a) a set of start:end token indices, b) a YES/NO answer if applicable (short answers ONLY), or c) a BLANK answer if no prediction can be made. The file should contain a header and have the following format:
+
+```
+-7853356005143141653_long,6:18
+-7853356005143141653_short,YES
+-545833482873225036_long,105:200
+-545833482873225036_short,
+-6998273848279890840_long,
+-6998273848279890840_short,NO
+```
+
+## Prizes
+
+link
+
+keyboard_arrow_up
+
+The following prizes will be awarded on the basis of private leaderboard rank.
+
+Leaderboard Prizes:
+
+- First Prize: $12,000
+
+- Second Prize: $8,000
+
+- Third Prize: $5,000
+
+* * * * *
+
+This competition is sponsored by the TensorFlow team. TensorFlow 2.0 was [recently released](https://medium.com/tensorflow/tensorflow-2-0-is-now-available-57d706c2a9ab) and this competition is to challenge Kagglers to use TensorFlow 2.0's APIs focused on usability, and easier, more intuitive development, to make advancements on the Natural Questions Document Understanding problem. Show off what you can do with TensorFlow 2.0 and you could win the following prizes.
+
+TensorFlow 2.0 Prizes:
+
+- First Prize: $12,000
+
+- Second Prize: $8,000
+
+- Third Prize: $5,000
+
+To be eligible for the TensorFlow 2.0 prizes, your code must use TensorFlow 2.0. More specifically:
+
+- it should run with a public release of TensorFlow 2.0 installed, and
+
+- should not use any tf.compat.v1 module symbols and should not raise deprecation warnings
+
+At the end of the competition, we'll publish a form for participants who have met these eligibility requirements to request consideration for these special set of TF2.0 prizes.
+
+## Timeline
+
+- January 15, 2020 - Team Merger deadline. This is the last day participants may join or merge teams.
+- January 15, 2020 - Entry deadline. You must accept the competition rules before this date in order to compete.
+- January 22, 2020 - Final submission deadline.
+
+All deadlines are at 11:59 PM UTC on the corresponding day unless otherwise noted. The competition organizers reserve the right to update the contest timeline if they deem it necessary.
+
+## Getting Started
+
+Here are some resources to help you get started:
+
+- [Starter notebook](https://www.kaggle.com/philculliton/using-tensorflow-2-0-w-bert-on-nq) using BERT on the dataset. *Note: This baseline uses code that was migrated from TF1.x. Be aware that it contains use of tf.compat.v1, which is not permitted to be eligible for [TF2.0 prizes in this competition](https://www.kaggle.com/c/tensorflow2-question-answering/overview/prizes). It is intended to be used as a starting point, but we're excited to see how much better you can using TF2.0!*
+- Some additional TensorFlow 2.0 resources:
+    - [Effective TensorFlow 2.0 Guide](https://www.tensorflow.org/guide/effective_tf2)
+    - [Migrating code from TF1.x to TF2.0](https://www.tensorflow.org/guide/migrate) as well as an [upgrade script](https://www.tensorflow.org/guide/upgrade)
+    - [tf.keras Guide](https://www.tensorflow.org/guide/keras/)
+    - [Kaggle's announcement about TF2.0 released in notebooks](https://www.kaggle.com/general/114059)
+- Information on the [TF2.0 prizes](https://www.kaggle.com/c/tensorflow2-question-answering/overview/prizes) you an earn in this competition
+- The [Natural Questions](https://ai.google.com/research/NaturalQuestions) site for more information about their dataset or to download directly from them. Note that this competition's private test set is completely different from any of the datasets used by Natural Questions.
+- Google Cloud Platform (GCP) Credits: For entrants in this competition, we are distributing $300 Coupons for GCP Credit. [Requests can be made by submitting this form](https://www.kaggle.com/GCP-Credits-Tensorflow2-QA) by November 18, 2019. We'll distribute the coupons by November 25th.
+    - You must have made a submission to this competition, or your request will be rejected.
+    - Only one person per team will be issued a Coupon. This will be verified and multiple requests from the same team will be de-duplicated.
+
+## Code Requirements
+
+![Kerneler](https://storage.googleapis.com/kaggle-media/competitions/general/Kerneler-white-desc2_transparent.png)
+
+### This is a code competition
+
+Submissions to this competition must be made through Kaggle Notebooks. In order for the "Submit to Competition" button to be active after a commit, the following conditions must be met:
+
+- GPU Notebooks <= 3 hours run-time
+- CPU Notebooks <= 9 hours run-time
+- Internet access off
+- External data is allowed (including training offline and uploading the model)
+- No custom packages enabled in kernels
+- Submission files must be named "submission.csv"
+
+Please see the [Code competition FAQ](https://www.kaggle.com/docs/competitions#kernels-only-FAQ) for more information on how to submit.
+
+## Citation
+
+keyboard_arrow_up
+
+fchollet, Julia Elliott, JWhyW24, Paige Bailey, Phil Culliton, raguptamd, TomK. (2019). TensorFlow 2.0 Question Answering. Kaggle. https://kaggle.com/competitions/tensorflow2-question-answering
+
+
+# Data
+
+## Dataset Description
+
+In this competition, we are tasked with selecting the best short and long answers from Wikipedia articles to the given questions.
+
+### What should I expect the data format to be?
+
+Each sample contains a Wikipedia article, a related question, and the candidate long form answers. The training examples also provide the correct long and short form answer or answers for the sample, if any exist.
+
+### What am I predicting?
+
+For each article + question pair, you must predict / select long and short form answers to the question drawn *directly from the article*.
+
+- A long answer would be a longer section of text that answers the question - several sentences or a paragraph.
+- A short answer might be a sentence or phrase, or even in some cases a YES/NO. The short answers are always contained within / a subset of one of the plausible long answers.
+- A given article can (and very often will) allow for both long *and* short answers, depending on the question.
+
+There is more detail about the data and what you're predicting [on the Github page for the Natural Questions dataset](https://github.com/google-research-datasets/natural-questions/blob/master/README.md). This page also contains helpful utilities and scripts. Note that we are using the simplified text version of the data - most of the HTML tags have been removed, and only those necessary to break up paragraphs / sections are included.
+
+### File descriptions
+
+- **simplified-nq-train.jsonl** - the training data, in newline-delimited JSON format.
+- **simplified-nq-kaggle-test.jsonl** - the test data, in newline-delimited JSON format.
+- **sample_submission.csv** - a sample submission file in the correct format
+
+### Data fields
+
+- **document_text** - the text of the article in question (with some HTML tags to provide document structure). The text can be tokenized by splitting on whitespace.
+- **question_text** - the question to be answered
+- **long_answer_candidates** - a JSON array containing all of the plausible long answers.
+- **annotations** - a JSON array containing all of the correct long + short answers. Only provided for train.
+- **document_url** - the URL for the full article. Provided for informational purposes only. This is NOT the simplified version of the article so indices from this cannot be used directly. The content may also no longer match the html used to generate document_text. Only provided for train.
+- **example_id** - unique ID for the sample.

dslighting/registry/tensorflow2-question-answering/description_obfuscated.md

@@ -0,0 +1,43 @@
+# Task
+
+For each article + question pair, you must predict / select long and short form answers to the question drawn *directly from the article*.
+
+- A long answer would be a longer section of text that answers the question - several sentences or a paragraph.
+- A short answer might be a sentence or phrase, or even in some cases a YES/NO. The short answers are always contained within / a subset of one of the plausible long answers.
+- A given article can (and very often will) allow for both long *and* short answers, depending on the question.
+
+There is more detail about the data and what you're predicting [on the Github page for the Natural Questions dataset](https://github.com/google-research-datasets/natural-questions/blob/master/README.md). This page also contains helpful utilities and scripts. Note that we are using the simplified text version of the data - most of the HTML tags have been removed, and only those necessary to break up paragraphs / sections are included.
+
+# Metric
+
+Micro F1. Predicted long and short answers must match exactly the token indices of one of the ground truth labels ((or match YES/NO if the question has a yes/no short answer). There may be up to five labels for long answers, and more for short. If no answer applies, leave the prediction blank/null.
+
+# Submission Format
+
+For each ID in the test set, you must predict a) a set of start:end token indices, b) a YES/NO answer if applicable (short answers ONLY), or c) a BLANK answer if no prediction can be made. The file should contain a header and have the following format:
+
+```
+-7853356005143141653_long,6:18
+-7853356005143141653_short,YES
+-545833482873225036_long,105:200
+-545833482873225036_short,
+-6998273848279890840_long,
+-6998273848279890840_short,NO
+```
+`
+# Data
+
+Each sample contains a Wikipedia article, a related question, and the candidate long form answers. The training examples also provide the correct long and short form answer or answers for the sample, if any exist.
+
+- **simplified-nq-train.jsonl** - the training data, in newline-delimited JSON format.
+- **simplified-nq-kaggle-test.jsonl** - the test data, in newline-delimited JSON format.
+- **sample_submission.csv** - a sample submission file in the correct format
+
+## Data fields
+
+- **document_text** - the text of the article in question (with some HTML tags to provide document structure). The text can be tokenized by splitting on whitespace.
+- **question_text** - the question to be answered
+- **long_answer_candidates** - a JSON array containing all of the plausible long answers.
+- **annotations** - a JSON array containing all of the correct long + short answers. Only provided for train.
+- **document_url** - the URL for the full article. Provided for informational purposes only. This is NOT the simplified version of the article so indices from this cannot be used directly. The content may also no longer match the html used to generate document_text. Only provided for train.
+- **example_id** - unique ID for the sample.

dslighting/registry/tensorflow2-question-answering/grade.py

@@ -0,0 +1,122 @@
+import pandas as pd
+from sklearn.metrics import f1_score
+
+from mlebench.grade_helpers import InvalidSubmissionError
+
+
+def prepare_for_metric(submission: pd.DataFrame, answers: dict):
+    """
+    `submission` is a pd.DataFrame with "example_id" and "PredictionString" columns.
+
+    `answers` is a dict in the format of "simplified-nq-train.jsonl" described here:
+    https://www.kaggle.com/c/tensorflow2-question-answering/data
+
+    The competition uses a micro F1 score metric, which is a binary classification metric.
+    For retrieval, we convert the submission and answers for each sample as:
+    - y_pred: if the submission has an answer, it is 1, else 0
+    - y_true:
+        - if the submission has an answer and it exists in the true labels, y_true=1
+        - if the submission has an answer but it does not exist in the true labels, y_true=0
+        - if the submission has no answer but there exists a true label, y_true=1
+        - if the submission has no answer and there is no true label, y_true=0
+
+    This is consistent with
+    https://www.kaggle.com/competitions/tensorflow2-question-answering/overview/evaluation
+    - TP = the predicted indices match one of the possible ground truth indices
+    - FP = the predicted indices do NOT match one of the possible ground truth indices, OR
+        a prediction has been made where no ground truth exists
+    - FN = no prediction has been made where a ground truth exists
+
+    Returns y_true, y_pred which are lists of 0s and 1s.
+    """
+    if len(submission) != 2 * len(answers):
+        raise InvalidSubmissionError(
+            f"Submission length {len(submission)} != 2 * answers length {len(answers)}"
+        )
+    # Empty strings are read as NaN by pandas, but we want these to remain empty strings
+    submission.fillna("", inplace=True)
+    submission = submission.astype(str)
+
+    expected_ids = []
+    for sample in answers:
+        expected_ids.append(f"{sample['example_id']}_long")
+        expected_ids.append(f"{sample['example_id']}_short")
+
+    if not ({"example_id", "PredictionString"}).issubset(submission.columns):
+        raise InvalidSubmissionError(
+            "Submission requires 'example_id' and 'PredictionString' columns."
+        )
+    if not set(submission["example_id"]) == set(expected_ids):
+        raise InvalidSubmissionError(
+            "Submission example_ids do not match expected_ids. Please ensure you have a "
+            "long and short answer for each example_id in the answers."
+        )
+    if not len(set(submission["example_id"])) == len(submission):
+        raise InvalidSubmissionError(
+            "Submission example_ids are not unique. Please ensure you have a "
+            "long and short answer for each example_id in the answers."
+        )
+
+    y_pred = [1 if row["PredictionString"].strip() else 0 for _, row in submission.iterrows()]
+
+    y_true = []
+    true_annotations = {sample["example_id"]: sample["annotations"][0] for sample in answers}
+    for idx, row in submission.iterrows():
+        sample_id = row["example_id"]
+
+        # Parse prediction
+        # Prediction may be any of "{start_token}:{end_token}", "YES", "NO", ""
+        pred = row["PredictionString"].strip()
+        if ":" in pred and len(pred.split(":")) == 2:
+            start_token, end_token = pred.split(":")
+            pred = (int(start_token), int(end_token))
+        elif pred in ["YES", "NO", ""]:
+            pass
+        else:
+            raise InvalidSubmissionError(f"Invalid submission format: {pred}")
+
+        # Parse valid answers
+        if sample_id.endswith("_long"):
+            sample_id = sample_id.split("_long")[0]
+            annotation = true_annotations[sample_id]
+            valid_answers = []
+            if annotation["long_answer"]["start_token"] != -1:
+                valid_answers.append(
+                    (
+                        int(annotation["long_answer"]["start_token"]),
+                        int(annotation["long_answer"]["end_token"]),
+                    ),
+                )
+        elif sample_id.endswith("_short"):
+            sample_id = sample_id.split("_short")[0]
+            annotation = true_annotations[sample_id]
+            valid_answers = [
+                (int(short_answer["start_token"]), int(short_answer["end_token"]))
+                for short_answer in annotation["short_answers"]
+            ]
+            if annotation["yes_no_answer"] != "NONE":
+                valid_answers.append(annotation["yes_no_answer"])
+        else:
+            raise InvalidSubmissionError(
+                f"sample_id {sample_id} does not end with '_long' or '_short'"
+            )
+
+        # Mark as 0 or 1 based on the above rules
+        if pred:
+            if pred in valid_answers:
+                y_true.append(1)
+            else:
+                y_true.append(0)
+        else:
+            if valid_answers:
+                y_true.append(1)
+            else:
+                y_true.append(0)
+
+    assert len(y_true) == len(y_pred), f"y_true {len(y_true)} != y_pred {len(y_pred)}"
+    return y_true, y_pred
+
+
+def grade(submission: pd.DataFrame, answers: dict) -> float:
+    y_true, y_pred = prepare_for_metric(submission, answers)
+    return f1_score(y_true=y_true, y_pred=y_pred, average="micro")

dslighting/registry/tensorflow2-question-answering/prepare.py

@@ -0,0 +1,122 @@
+from pathlib import Path
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from tqdm import tqdm
+
+from mlebench.utils import get_logger
+
+logger = get_logger(__name__)
+
+
+def extract_string(document_text: str, start_token_idx: int, end_token_idx: int) -> str:
+    document_tokens = document_text.split(" ")
+    extract_tokens = document_tokens[start_token_idx:end_token_idx]
+    return " ".join(extract_tokens)
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Splits the data in raw into public and private datasets with appropriate test/train splits.
+    """
+
+    # Create train, test from train split
+    train_file = "simplified-nq-train.jsonl"
+
+    logger.info("Counting lines in train file...")
+    with open(raw / train_file, "r") as f:
+        n_lines = sum(1 for _ in f)
+    logger.info(f"Found {n_lines} lines in train file.")
+
+    # Read data in chunks to avoid memory issues
+    train_ids, test_ids = [], []
+    lightweight_test = []  # We'll use this to create a gold submission later
+    with tqdm(total=n_lines, desc="Splitting data") as pbar:
+        for df in pd.read_json(raw / train_file, orient="records", lines=True, chunksize=1_000):
+            # Convert IDs to strings, Kaggle.com is inconsistent about this but strings make more sense
+            df["example_id"] = df["example_id"].astype(str)
+            new_train, new_test = train_test_split(df, test_size=0.1, random_state=0)
+
+            keys_to_keep = [
+                "example_id",
+                "question_text",
+                "document_text",
+                "long_answer_candidates",
+            ]
+            new_test_without_labels = new_test.copy()[keys_to_keep]
+
+            # Append lines to new train and test
+            with open(public / "simplified-nq-train.jsonl", "a") as f:
+                f.write(new_train.to_json(orient="records", lines=True))
+            with open(private / "test.jsonl", "a") as f:
+                f.write(new_test.to_json(orient="records", lines=True))
+            with open(public / "simplified-nq-test.jsonl", "a") as f:
+                f.write(new_test_without_labels.to_json(orient="records", lines=True))
+
+            train_ids.extend(new_train["example_id"].tolist())
+            test_ids.extend(new_test["example_id"].tolist())
+            lightweight_test.append(
+                new_test.copy()[["example_id", "question_text", "annotations"]]
+            )  # For gold submission
+            pbar.update(len(df))
+
+    lightweight_test = pd.concat(lightweight_test, ignore_index=True)
+
+    assert len(train_ids) + len(test_ids) == n_lines
+    assert len(lightweight_test) == len(test_ids)
+
+    # Create a gold submission with columns "example_id", "PredictionString"
+    gold_rows = []
+    for idx, sample in tqdm(
+        lightweight_test.iterrows(), total=len(lightweight_test), desc="Creating gold submission"
+    ):
+        sample = sample.to_dict()
+        assert len(sample["annotations"]) == 1
+        annotation = sample["annotations"][0]
+
+        # Create short answer
+
+        # Multiple answers are possible: yes_no_answer or one of short_answers
+        # We just take the first one
+        if annotation["yes_no_answer"] != "NONE":
+            answer = annotation["yes_no_answer"]
+        elif len(annotation["short_answers"]) > 0:
+            start_token = annotation["short_answers"][0]["start_token"]
+            end_token = annotation["short_answers"][0]["end_token"]
+            answer = f"{start_token}:{end_token}"
+        else:
+            answer = ""
+
+        logger.debug(f"q: {sample['question_text']}")
+        logger.debug(f"a: {answer}")
+        logger.debug("")
+
+        gold_rows.append(
+            {"example_id": f"{sample['example_id']}_short", "PredictionString": answer}
+        )
+
+        # Create long answer
+
+        if annotation["long_answer"]["start_token"] != -1:
+            start_token = annotation["long_answer"]["start_token"]
+            end_token = annotation["long_answer"]["end_token"]
+            answer = f"{start_token}:{end_token}"
+        else:
+            answer = ""
+
+        logger.debug(f"q: {sample['question_text']}")
+        logger.debug(f"a: {answer}")
+        logger.debug("")
+
+        gold_rows.append({"example_id": f"{sample['example_id']}_long", "PredictionString": answer})
+
+    gold_submission = pd.DataFrame(gold_rows)
+    gold_submission.to_csv(private / "gold_submission.csv", index=False)
+
+    # Sample submission
+    sample_submission = gold_submission.copy()
+    sample_submission["PredictionString"] = ""
+    sample_submission.to_csv(public / "sample_submission.csv", index=False)
+
+    assert len(gold_submission) == 2 * len(test_ids)
+    assert len(sample_submission) == 2 * len(test_ids)