code-searcher/cosqa_data_code.py at main · rlucas7/code-searcher · 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
""" This module houses code to generate relevances for the cosQA data.
    This data serves as a comparison and benchmark for our works.
    This module is a 'data shim', it sits between the format of cosQA that is publicly available
    and the format we use in this project for annotations of code search.
    Once this module is executed, the expected output is a csv file whose
    subsequent use should work 'as-is' as an input for the `gen-llm-rels` click command.
    Then once those relevances are determined we can use them for metrics calculation in the
    `gen-ir-metrics` click command. Therefore, we use this module only once in the workflow for
    comparing the results. The paper that describes the cosQA data is here: https://arxiv.org/abs/2105.13239
"""

from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from csv import DictWriter
from json import loads
from random import randint

from requests import get as RequestsGet

parser = ArgumentParser(
    prog='CosQA data shim script',
    formatter_class=ArgumentDefaultsHelpFormatter,
    description="A script to convert the cosQA data into a format useable in our click command workflows."
    )

data_url = "https://raw.githubusercontent.com/Jun-jie-Huang/CoCLR/refs/heads/main/data/qa/cosqa-dev.json"

parser.add_argument("--data_url", type=str, help="The url to access the data in json form.", default=data_url)
parser.add_argument("--output_filename", type=str, help="The name of the output file generated by this script.", default="rad.csv")


if __name__ == "__main__":
    args = parser.parse_args()
    ## NOTE: this is not the full dataset of 20K+, it's small and easy to work with though
    resp = RequestsGet(args.data_url)
    dev_contents = loads(resp.content)

    # a couple quick spot checks, this is mostly to if the data format changed or there was a transmission error

    N = 10 # number to check
    for i in range(N):
        random_idx = randint(0, len(dev_contents))
        assert int(dev_contents[random_idx]['idx'].split("-")[-1]) == random_idx, f"error index not expected value at: {random_idx}..."


    # each indexed entry has keys:
    # i. 'idx' a string with value f"cosqa-dev-{idx}", where `idx` starts from 0
    # ii. 'doc' the plaintext search query into bing search engine
    # iii. 'code' some python code from github
    # iv. 'code_tokens' tokenized code
    # v. 'docstring_tokens': tokenized docstring for the python code returned
    # vi. 'label': a 0/1 value, 1 if the code is relevant to the search


    ## Now we make this fit into the format expected for the `gen-llm-rels` command in db.py module
    ## basically we need a file with a header row, and then the rows themselves with some data...
    # we'll make filler data for things that aren't used for the relevance determination by the LLMs

    with open(args.output_filename, "w", newline="\n") as csv_file:
        fieldnames = [
            "query_id",
            "post_id",
            "user_id",
            "relevances",
            "rank",
            "distance",
            "query",
            "doc",
            "code",
        ]
        dw = DictWriter(
            csv_file,
            delimiter="|",
            quotechar='"',
                    fieldnames=fieldnames,
                    lineterminator="\r\n",
        )
        dw.writeheader()
        for row in enumerate(dev_contents, start=1):
            code_n_docstring = row[1]['code'].split('"""')
            docstring_only  = '"'*1 + code_n_docstring[1].rstrip(" ") + '"'*1
            code_only = code_n_docstring[0] + code_n_docstring[2]
            datum = {
                'query_id': row[1]['idx'],
                'post_id': row[1]['idx'],
                'user_id': 77,
                'relevances': row[1]['label'],
                'rank': 1,
                'distance': 1,
                'query': row[1]['doc'], # confusing name but this *is* the query
                'doc': docstring_only,
                'code': code_only
            }
            ## soe formatting to the row...
            dw.writerow(datum)

    ## now you can test reading the output file using the click command `gen-llm-rels`, something
    ## like the following:
    ##
    ## ```bash
    ## flask --app vec_search gen-llm-rels rad.csv llm_gen_rel.csv openai
    ## ````
    ## once you run this script.