Spaces:
Runtime error
Runtime error
File size: 1,771 Bytes
e9097df 019cb59 e9097df 019cb59 e9097df 030f79a e9097df 030f79a e9097df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
# -*- coding: utf-8 -*-
"""code-search.ipynb
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/drive/1-TlihNx5XCiVSxUHDF1oHFNcfpuy_k0N
"""
# Install Cohere for embeddings
import cohere
import numpy as np
import pandas as pd
import gradio as gr
import os
from sklearn.metrics.pairwise import cosine_similarity
from annoy import AnnoyIndex
import warnings
warnings.filterwarnings('ignore')
pd.set_option('display.max_colwidth', None)
data_df = pd.read_csv('functions_data.csv')
#data_df.head()
data_df['docstring'].fillna('not specified', inplace=True)
# Paste your API key here. Remember to not share publicly
key = os.environ.get('API_KEY')
api_key = key
# Create and retrieve a Cohere API key from dashboard.cohere.ai/welcome/register
co = cohere.Client(api_key)
search_index = AnnoyIndex(4096, 'angular')
search_index.load('code.ann') # super fast, will just mmap the file
def get_code(query):
# Get the query's embedding
query_embed = co.embed(texts=[query],
model="large",
truncate="LEFT").embeddings
# Retrieve the nearest neighbors
similar_item_ids = search_index.get_nns_by_vector(query_embed[0],1)
return data_df.iloc[similar_item_ids[0]]['function_body'], data_df.iloc[similar_item_ids[0]]['file_path']
examples = ['compute diffusion of given data']
inputs = gr.Textbox(label='query')
outputs = [gr.Textbox(label='matched function'), gr.Textbox(label='File path')
title = "Search Code"
description = "Semantically search codebase using Cohere embed API. This demo uses Open AI point cloud codebase as an example"
iface = gr.Interface(fn=get_code, inputs=inputs, outputs=outputs, description = description)
iface.launch()
|