Spaces:
Runtime error
Runtime error
File size: 2,690 Bytes
d6585f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
#
# Pyserini: Reproducible IR research with sparse and dense representations
#
# 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.
#
import argparse
import os
import csv
import pyserini.collection
def escape_title(title):
# Neo4j CSV loader is sensitive to double qoutes
return title.replace("\"", "\\'").replace("\\", "")
def main(path):
collection = pyserini.collection.Collection('Cord19AbstractCollection', path)
articles = collection.__next__()
with open("articles.csv", 'w') as article_csv, open("edges.csv", 'w') as edge_csv:
article_csv = csv.writer(article_csv)
edge_csv = csv.writer(edge_csv)
article_csv.writerow(["cord_uid", "title", "pmcid"])
edge_csv.writerow(["cord_uid", "target_title", "doi"])
prev_titles = set()
prev_cord_uid = set()
for d in articles:
article = pyserini.collection.Cord19Article(d.raw)
title = article.title()
cord_uid = article.cord_uid()
if article.is_full_text() and title and title not in prev_titles \
and cord_uid not in prev_cord_uid:
article_data = [article.cord_uid(), escape_title(title),
article.json["paper_id"]]
article_csv.writerow(article_data)
prev_titles.add(title)
prev_cord_uid.add(cord_uid)
bib_entries = article.bib_entries()
# Create edge between article and each cited title
for bib_ref in bib_entries:
ref = bib_entries[bib_ref]
if ref['title']:
doi = ref['other_ids'].get('DOI')
doi = None if doi == [] or doi is None else doi[0]
edge = [article.cord_uid(), escape_title(ref['title']), doi]
edge_csv.writerow(edge)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Load CORD-19 citation data into Neo4j')
parser.add_argument('--path', type=str, required=True,
help='The path to CORD-19 collection')
args = parser.parse_args()
main(args.path)
|