<jupyter_start><jupyter_text>Code searchWe index our own openai-python code repository, and show how it can be searched. We implement a simple version of file parsing and extracting of functions from python files. The dataset is created in the [Obtain_dataset Notebook](Obtain_dataset.ipynb).<jupyter_code>import os
from glob import glob
import pandas as pd

def get_function_name(code):
    """
    Extract function name from a line beginning with "def "
    """
    assert code.startswith("def ")
    return code[len("def "): code.index("(")]

def get_until_no_space(all_lines, i) -> str:
    """
    Get all lines until a line outside the function definition is found.
    """
    ret = [all_lines[i]]
    for j in range(i + 1, i + 10000):
        if j < len(all_lines):
            if len(all_lines[j]) == 0 or all_lines[j][0] in [" ", "\t", ")"]:
                ret.append(all_lines[j])
            else:
                break
    return "\n".join(ret)

def get_functions(filepath):
    """
    Get all functions in a Python file.
    """
    whole_code = open(filepath).read().replace("\r", "\n")
    all_lines = whole_code.split("\n")
    for i, l in enumerate(all_lines):
        if l.startswith("def "):
            code = get_until_no_space(all_lines, i)
            function_name = get_function_name(code)
            yield {"code": code, "function_name": function_name, "filepath": filepath}


# get user root directory
root_dir = os.path.expanduser("~")

# path to code repository directory
code_root = root_dir + "/openai-python"
code_files = [y for x in os.walk(code_root) for y in glob(os.path.join(x[0], '*.py'))]
print("Total number of py files:", len(code_files))
all_funcs = []
for code_file in code_files:
    funcs = list(get_functions(code_file))
    for func in funcs:
        all_funcs.append(func)

print("Total number of functions extracted:", len(all_funcs))<jupyter_output>Total number of py files: 40
Total number of functions extracted: 64<jupyter_text>For code search models we use babbage-code-search-code to obtain embeddings for code snippets, and code-search-text to embed natural language queries.<jupyter_code>from utils import get_embedding

df = pd.DataFrame(all_funcs)
df['code_embedding'] = df['code'].apply(lambda x: get_embedding(x, engine='babbage-code-search-code'))
df['filepath'] = df['filepath'].apply(lambda x: x.replace(code_root, ""))
df.to_csv("output/code_search_openai-python.csv", index=False)
df.head()
from utils import cosine_similarity

def search_functions(df, code_query, n=3, pprint=True, n_lines=7):
    embedding = get_embedding(code_query, engine='babbage-code-search-text')
    df['similarities'] = df.code_embedding.apply(lambda x: cosine_similarity(x, embedding))

    res = df.sort_values('similarities', ascending=False).head(n)
    if pprint:
        for r in res.iterrows():
            print(r[1].filepath+":"+r[1].function_name + "  score=" + str(round(r[1].similarities, 3)))
            print("\n".join(r[1].code.split("\n")[:n_lines]))
            print('-'*70)
    return res
res = search_functions(df, 'Completions API tests', n=3)

res = search_functions(df, 'fine-tuning input data validation logic', n=3)
res = search_functions(df, 'find common suffix', n=2, n_lines=10)
res = search_functions(df, 'Command line interface for fine-tuning', n=1, n_lines=20)<jupyter_output>/openai/cli.py:tools_register  score=0.651
def tools_register(parser):
    subparsers = parser.add_subparsers(
        title="Tools", help="Convenience client side tools"
    )

    def help(args):
        parser.print_help()

    parser.set_defaults(func=help)

    sub = subparsers.add_parser("fine_tunes.prepare_data")
    sub.add_argument(
        "-f",
        "--file",
        required=True,
        help="JSONL, JSON, CSV, TSV, TXT or XLSX file containing prompt-completion examples to be analyzed."
        "This should be the local file path.",
    )
    sub.add_argument(
        "-q",
----------------------------------------------------------------------
{'fraction_non_alphanumeric': 0.1464081531195625, 'fraction_numerical': 0.008451404424558787, 'mean_word_length': 3.402625820568928, 'pattern_counts': {'":': 4, '<': 8, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '190773', 'license': 'MIT', 'n_tokens_mistral': 1319, 'n_tokens_neox': 1282, 'n_words': 340, 'path': 'examples/embeddings/Code_search.ipynb', 'repo_name': 'BastianZim/openai-python', 'size': '4023'}