Spaces:
Runtime error
Runtime error
| # | |
| # 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) | |