AISafetyGraph/generate_md/explore_ds.py at main · ai-safety-graph/AISafetyGraph · GitHub

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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#%%
# read the file dataset/alignment_texts.jsonl
import jsonlines
from pprint import pprint

dataset_file = "dataset/alignment_texts.jsonl"

dictionary = {}
with jsonlines.open("dataset/alignment_texts.jsonl", "r") as reader:
  for entry in reader:
    # print(type(entry)) # dict
    # print(entry.keys())
    # dict_keys(['title', 'title_detail',
    # 'links', 'link',
    # 'authors', 'author', 'author_detail',
    # 'published', 'published_parsed',
    # 'tags', 'id', 'guidislink',
    # 'summary', 'summary_detail',
    # 'content',
    # 'post-id', 'text',
    # 'source', 'source_type'])

    try:
      if not entry["source"] == "arxiv":
        continue
    except KeyError:
      continue

    print(f"{entry['source']=}")
    # print(entry.keys())
    # dict_keys([
    # 'source', 'source_type',
    # 'converted_with', 'paper_version',
    # 'title', 'authors',
    # 'date_published', 'data_last_modified', 'url',
    # 'abstract',
    # 'author_comment', 'journal_ref', 'doi',
    # 'primary_category', 'categories',
    # 'citation_level',
    # 'alignment_text', 'confidence_score',
    # 'main_tex_filename', 'text',
    # 'bibliography_bbl', 'bibliography_bib',
    # 'arxiv_citations'])

    try:
      print(f"{entry["title"]=}")
      # print(f"{entry["primary_category"]=}")
      # print(f"{entry["categories"]=}")
      # print(f"{entry["alignment_text"]=}")
      # print(f"{entry["confidence_score"]=}")
      # print(f"{entry["text"]=}")
      print(f"{entry["abstract"]=}")
    except KeyError:
      print("title key not found")
      continue

    break

    # write all to a yaml file
    # break
# %%
def extract_arxiv_papers_to_md(output_file):
    input_file = dataset_file
    with open(output_file, "w", encoding="utf-8") as out_file:
        paper_count = 0

        with jsonlines.open(input_file, "r") as reader:
            for entry in reader:
                try:
                    if entry["source"] != "arxiv":
                        continue

                    title = entry["title"]
                    abstract = entry["abstract"]

                    # Write the paper information in Markdown format
                    out_file.write(f"# {title}\n\n")
                    out_file.write(f"{abstract}\n\n")
                    out_file.write("---\n\n")  # Separator between papers

                    paper_count += 1

                    if paper_count >= 20:
                       break

                except KeyError:
                    continue

    print(f"Extracted {paper_count} arXiv papers to {output_file}")

# Usage
# output_file = "arxiv_papers_for_llm.md"
# extract_arxiv_papers_to_md(output_file)

def extract_arxiv_papers_to_yaml(output_file):
    # extract title, abstract and url
    input_file = dataset_file
    with open(output_file, "w", encoding="utf-8") as out_file:
        paper_count = 0

        with jsonlines.open(input_file, "r") as reader:
            for entry in reader:
                try:
                    if entry["source"] != "arxiv":
                        continue

                    title = entry["title"]
                    abstract = entry["abstract"]
                    url = entry["url"]

                    # remove colon from title and abstract
                    title = title.replace(":", "")
                    abstract = abstract.replace(":", "")
                    title = title.replace("\"", "")
                    abstract = abstract.replace("\"", "")

                    title = title.replace("`", "")
                    abstract = abstract.replace("`", "")

                    # Write the paper information in YAML format
                    out_file.write(f"- title: {title}\n")
                    out_file.write(f"  abstract: {abstract}\n")
                    out_file.write(f"  url: {url}\n\n")

                    paper_count += 1

                    # if paper_count >= 20:
                    #    break

                except KeyError:
                    continue

# Usage
output_file = "arxiv_papers_for_llm.yaml"
extract_arxiv_papers_to_yaml(output_file)