{"cells":[{"metadata":{},"cell_type":"markdown","source":"# TREC-COVID Submission\n\nThis notebook builds a submission file using the search index build in the [TREC-COVID Search Index](https://www.kaggle.com/davidmezzetti/trec-covid-search-index) notebook.\n\nFor each topic, a query is run against the search index and the Top N search results are saved. Upon completion, a file with the search results per topic are written to an output file named submission.csv","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"# Install cord19q project\n!pip install git+https://github.com/neuml/cord19q","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"import csv\nimport os\nimport shutil\n\nimport pandas as pd\n\nfrom cord19q.models import Models\nfrom cord19q.query import Query\n\n# Workaround for mdv terminal width issue\nos.environ[\"COLUMNS\"] = \"80\"\n\ndef uids():\n    # Entry date mapping sha id to date\n    uids = {}\n    \n    # Load in memory date lookup\n    with open(\"../input/cord-19-article-entry-dates/entry-dates.csv\", mode=\"r\") as csvfile:\n        for row in csv.DictReader(csvfile):\n            uids[row[\"sha\"]] = row[\"cord_uid\"]\n\n    return uids\n\n# Copy vectors locally for predictable performance\nshutil.copy(\"../input/cord19-fasttext-vectors/cord19-300d.magnitude\", \"/tmp\")\n\n# Load sha - cord id mapping\nidmap = uids()\n\nsubmission = []\ntopn = 50\n\n# Load model\nembeddings, db = Models.load(\"../input/trec-covid-search-index/cord19q\")\ncur = db.cursor()\n\nwith open(\"../input/trec-covid-information-retrieval/topics-rnd3.csv\", mode=\"r\") as csvfile:\n    for topic in csv.DictReader(csvfile):\n        # Run the search\n        results = Query.search(embeddings, cur, topic[\"query\"], topn)\n        \n        # Get results grouped by document\n        documents = Query.documents(results, topn)   \n\n        for uid in documents:\n            # uid is third element, lookup cord_uid from shas:\n            submission.append((topic[\"topic-id\"], idmap[uid]))\n            \ndf = pd.DataFrame(submission, columns=[\"topic-id\", \"cord-id\"])\ndf.to_csv(\"submission.csv\", index=False)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}