openai的API实现代码函数检索

openwin_top

已于 2024-03-09 20:32:52 修改

阅读量460

点赞数

CC 4.0 BY-SA版权

分类专栏：人工智能文章标签： python 机器学习深度学习

于 2023-03-19 08:55:46 首次发布

本文链接：https://blog.youkuaiyun.com/zhangzhechun/article/details/129646565

人工智能专栏收录该内容

55 篇文章

订阅专栏

该代码段用于从Python代码文件中提取函数，利用OpenAI的预训练模型获取函数的向量表示，然后通过余弦相似性进行搜索。支持对大量Python代码文件进行处理，找到与查询函数内容相似的函数。同时，文章提到了opencv库在图像处理，如物体跟踪、人脸识别等方面的运用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

从代码库中先把所有的函数提取出来

import os
from glob import glob
import pandas as pd

def get_function_name(code):
    assert code.startswith("def ")
    return code[len("def "): code.index("(")]

def get_until_no_space(all_lines, i) -> str:
    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):
    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}


root_dir = os.path.expanduser("~")

code_root = root_dir + "/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))

if len(code_files) == 0:
    print("no 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))

OpenAI API中的get_embedding是用于获取给定文本的向量表示的方法。它使用OpenAI的预训练模型，将输入文本转换为一个向量，该向量可用于计算文本之间的相似性或进行其他自然语言处理任务。
get_embedding方法接受一个文本字符串作为输入，并返回一个表示该文本的向量。这个向量通常是一个具有固定维度的浮点数数组。该方法可以用于多种应用，例如文本分类、聚类、信息检索、自动摘要、机器翻译等。
需要注意的是，get_embedding方法返回的向量是基于OpenAI的预训练模型计算得出的，并且可能不适用于某些特定应用场景，因此需要根据具体应用情况选择合适的方法和模型。

from openai.embeddings_utils import get_embedding

df = pd.DataFrame(all_funcs)
df['code_embedding'] = df['code'].apply(lambda x: get_embedding(x, engine='text-embedding-ada-002'))
df['filepath'] = df['filepath'].apply(lambda x: x.replace(code_root, ""))
df.to_csv("data/code_search_python.csv", index=False)
df.head()

定义搜索函数

from openai.embeddings_utils import cosine_similarity

def search_functions(df, code_query, n=3, pprint=True, n_lines=7):
    embedding = get_embedding(code_query, engine='text-embedding-ada-002')
    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, '想要搜索的函数内容', n=3)