Shamima/code-search

1# -*- coding: utf-8 -*-2 3# Install Cohere for embeddings4 5import cohere6import numpy as np7import pandas as pd8import gradio as gr9import os10from sklearn.metrics.pairwise import cosine_similarity11from annoy import AnnoyIndex12import warnings13warnings.filterwarnings('ignore')14pd.set_option('display.max_colwidth', None)15 16data_df = pd.read_csv('functions_data.csv')17 18 19data_df['docstring'].fillna('not specified', inplace=True)20 21# Paste your API key here. Remember to not share publicly22key = os.environ.get('API_KEY')23api_key = key24 25# Create and retrieve a Cohere API key from dashboard.cohere.ai/welcome/register26co = cohere.Client(api_key)27 28 29search_index = AnnoyIndex(4096, 'angular')30search_index.load('code.ann') # super fast, will just mmap the file31 32def get_code(query):33    # Get the query's embedding34    query_embed = co.embed(texts=[query],35                    model="large",36                    truncate="LEFT").embeddings37 38    # Retrieve the nearest neighbors39    similar_item_ids = search_index.get_nns_by_vector(query_embed[0],1)40    41   42    return data_df.iloc[similar_item_ids[0]]['function_body'], data_df.iloc[similar_item_ids[0]]['file_path']43examples = ['compute diffusion of given data']44inputs = gr.Textbox(label='query')45outputs = [gr.Textbox(label='matched function'), gr.Textbox(label='File path')]46title = "Search Code"47description = "Semantically search codebase using Cohere embed API. This demo uses Open AI point cloud codebase https://github.com/openai/point-e as an example"48iface = gr.Interface(fn=get_code, inputs=inputs, outputs=outputs, description = description, examples=examples, title=title)49iface.launch()50 51 52 53