جستجوی کد با استفاده از embeddings

 1import pandas as pd
 2from pathlib import Path
 3
 4DEF_PREFIXES = ['def ', 'async def ']
 5NEWLINE = '\n'
 6
 7def get_function_name(code):
 8    """
 9    Extract function name from a line beginning with 'def' or 'async def'.
10    """
11    for prefix in DEF_PREFIXES:
12        if code.startswith(prefix):
13            return code[len(prefix): code.index('(')]
14
15
16def get_until_no_space(all_lines, i):
17    """
18    Get all lines until a line outside the function definition is found.
19    """
20    ret = [all_lines[i]]
21    for j in range(i + 1, len(all_lines)):
22        if len(all_lines[j]) == 0 or all_lines[j][0] in [' ', '\t', ')']:
23            ret.append(all_lines[j])
24        else:
25            break
26    return NEWLINE.join(ret)
27
28
29def get_functions(filepath):
30    """
31    Get all functions in a Python file.
32    """
33    with open(filepath, 'r') as file:
34        all_lines = file.read().replace('\r', NEWLINE).split(NEWLINE)
35        for i, l in enumerate(all_lines):
36            for prefix in DEF_PREFIXES:
37                if l.startswith(prefix):
38                    code = get_until_no_space(all_lines, i)
39                    function_name = get_function_name(code)
40                    yield {
41                        'code': code,
42                        'function_name': function_name,
43                        'filepath': filepath,
44                    }
45                    break
46
47
48def extract_functions_from_repo(code_root):
49    """
50    Extract all .py functions from the repository.
51    """
52    code_files = list(code_root.glob('**/*.py'))
53
54    num_files = len(code_files)
55    print(f'Total number of .py files: {num_files}')
56
57    if num_files == 0:
58        print('Verify openai-python repo exists and code_root is set correctly.')
59        return None
60
61    all_funcs = [
62        func
63        for code_file in code_files
64        for func in get_functions(str(code_file))
65    ]
66
67    num_funcs = len(all_funcs)
68    print(f'Total number of functions extracted: {num_funcs}')
69
70    return all_funcs

 1
 2# Set user root directory to the 'openai-python' repository
 3root_dir = Path.home()
 4
 5# Clone the repo
 6git clone https://github.com/openai/openai-python.git
 7
 8code_root = root_dir / 'openai-python'
 9
10# Extract all functions from the repository
11all_funcs = extract_functions_from_repo(code_root)

 1from openai import OpenAI # for calling the OpenAI API
 2import os
 3
 4client = OpenAI(
 5    api_key=os.environ.get(("GILAS_API_KEY", "<کلید API خود را اینجا بسازید https://dashboard.gilas.io/apiKey>")), 
 6    base_url="https://api.gilas.io/v1/" # Gilas APIs
 7)
 8
 9def get_embedding(query)
10    query_embedding_response = client.embeddings.create(
11        model=embedding_model,
12        input=query,
13    )
14    return query_embedding_response.data[0].embedding
15
16df = pd.DataFrame(all_funcs)
17df['code_embedding'] = df['code'].apply(lambda x: get_embedding(x, model='text-embedding-3-small'))
18df['filepath'] = df['filepath'].map(lambda x: Path(x).relative_to(code_root))
19df.to_csv("data/code_search_openai-python.csv", index=False)
20df.head()

 1def cosine_similarity(vec1, vec2):
 2    """Calculate the cosine similarity between two vectors."""
 3    vec1 = np.array(vec1, dtype=float)
 4    vec2 = np.array(vec2, dtype=float)
 5
 6
 7    dot_product = np.dot(vec1, vec2)
 8    norm_vec1 = np.linalg.norm(vec1)
 9    norm_vec2 = np.linalg.norm(vec2)
10    return dot_product / (norm_vec1 * norm_vec2)
11
12
13def search_functions(df, code_query, n=3, pprint=True, n_lines=7):
14    embedding = get_embedding(code_query, model='text-embedding-3-small')
15    df['similarities'] = df.code_embedding.apply(lambda x: cosine_similarity(x, embedding))
16
17    res = df.sort_values('similarities', ascending=False).head(n)
18
19    if pprint:
20        for r in res.iterrows():
21            print(f"{r[1].filepath}:{r[1].function_name}  score={round(r[1].similarities, 3)}")
22            print("\n".join(r[1].code.split("\n")[:n_lines]))
23            print('-' * 70)
24
25    return res

1res = search_functions(df, 'fine-tuning input data validation logic', n=3)

1res = search_functions(df, 'find common suffix', n=2, n_lines=10)

1res = search_functions(df, 'Command line interface for fine-tuning', n=1, n_lines=20)