# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# © 2020 Nokia
#
# Licensed under the BSD 3 Clause license
#
# SPDX-License-Identifier: BSD-3-Clause

# # Prepare Conala snippet collection and evaluation data

# +
from pathlib import Path
import json
from collections import defaultdict

from codesearch.data import load_jsonl, save_jsonl

corpus_url = "http://www.phontron.com/download/conala-corpus-v1.1.zip"
conala_dir = Path("conala-corpus")
conala_train_fn = conala_dir/"conala-test.json"
conala_test_fn = conala_dir/"conala-train.json"
conala_mined_fn = conala_dir/"conala-mined.jsonl"

conala_snippets_fn = "conala-curated-snippets.jsonl"
conala_retrieval_test_fn = "conala-test-curated-0.5.jsonl"

if not conala_train_fn.exists():
    # !wget $corpus_url
    # !unzip conala-corpus-v1.1.zip
# -

conala_mined = load_jsonl(conala_mined_fn)

# The mined dataset seems to noisy to incorporate in the snippet collection:

# + tags=[]
# !sed -n '10000,10009p;10010q' $conala_mined_fn

# +
with open(conala_train_fn) as f:
    conala_train = json.load(f)

with open(conala_test_fn) as f:
    conala_test = json.load(f)

conala_all = conala_train + conala_test
conala_all[:2], len(conala_all), len(conala_train), len(conala_test)

# + tags=[]
for s in conala_all:
    if s["rewritten_intent"] == "Convert the first row of numpy matrix `a` to a list":
        print(s)

# +
question_ids = {r["question_id"] for r in conala_all}
intents = set(r["intent"] for r in conala_all)

len(question_ids), len(conala_all), len(intents)
# -

id2snippet = defaultdict(list)
for r in conala_all:
    id2snippet[r["question_id"]].append(r)

for r in conala_all:
    if not r["intent"]: 
        print(r)
    if r["intent"].lower() == (r["rewritten_intent"] or "").lower():
        print(r)

# +
import random
random.seed(42)

snippets = []
eval_records = []
for question_id in id2snippet:
    snippets_ = [r for r in id2snippet[question_id] if r["rewritten_intent"]]
    if not snippets_: continue
    for i, record in enumerate(snippets_):

        snippet_record = {
            "id": f'{record["question_id"]}-{i}', 
            "code": record["snippet"], 
            "description": record["rewritten_intent"],
            "language": "python",
            "attribution": f"https://stackoverflow.com/questions/{record['question_id']}"
        }
        snippets.append(snippet_record)
    # occasionally snippets from the same question have a slightly different intent
    # to avoid similar queries, we create only one query per question
    query = random.choice(snippets_)["intent"]
    
    if any(query.lower() == r["description"].lower() for r in snippets[-len(snippets_):] ):
        print(f"filtering query {query}")
        continue
    relevant_ids = [r["id"] for r in snippets[-len(snippets_):] ]
    
    eval_records.append({"query": query, "relevant_ids": relevant_ids})
    
snippets[:2], len(snippets), eval_records[:2], len(eval_records)
# -

id2snippet_ = {r["id"]: r for r in snippets}

for i, eval_record in enumerate(eval_records):
    print(f"Query: {eval_record['query']}")
    print(f"Relevant descriptions: {[id2snippet_[id]['description'] for id in eval_record['relevant_ids']]}")
    if i == 10:
        break

# + tags=[]
from codesearch.text_preprocessing import compute_overlap

compute_overlap("this is a test", "test test")
# -

overlaps = []
filtered_eval_records = []
for r in eval_records:
    query = r["query"]
    descriptions = [id2snippet_[id]['description'] for id in r['relevant_ids']]
    overlap = max(compute_overlap(query, d)[1] for d in descriptions)
    overlaps.append(overlap)
    
    if overlap < 0.5 :
        filtered_eval_records.append(r)
filtered_eval_records[:2], len(filtered_eval_records)

save_jsonl(conala_snippets_fn, snippets)
save_jsonl(conala_retrieval_test_fn, filtered_eval_records)

StarcoderdataJupyterScriptsDedupFiltered