startup-finder / app.py
Michelangiolo's picture
changes
3d4e8d2
import os
os.system('pip install openpyxl')
os.system('pip install sentence-transformers')
import pandas as pd
import gradio as gr
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('all-mpnet-base-v2') #all-MiniLM-L6-v2 #all-mpnet-base-v2
df = pd.read_parquet('df_encoded3.parquet')
df['tags'] = df['tags'].apply(lambda x : str(x))
def parse_raised(x):
if x == 'Undisclosed':
return 0
else:
quantifier = x[-1]
x = float(x[1:-1])
if quantifier == 'K':
return x/1000
elif quantifier == 'M':
return x
df['raised'] = df['raised'].apply(lambda x : parse_raised(x))
df['stage'] = df['stage'].apply(lambda x : x.lower())
df = df.reset_index(drop=True)
from sklearn.neighbors import NearestNeighbors
import pandas as pd
from sentence_transformers import SentenceTransformer
nbrs = NearestNeighbors(n_neighbors=5000, algorithm='ball_tree').fit(df['text_vector_'].values.tolist())
def search(df, query):
product = model.encode(query).tolist()
# product = df.iloc[0]['text_vector_'] #use one of the products as sample
#prepare model
#
distances, indices = nbrs.kneighbors([product]) #input the vector of the reference object
#print out the description of every recommended product
return df.iloc[list(indices)[0]][['name', 'raised', 'target', 'size', 'stage', 'country', 'source', 'description', 'tags']]
def filter_df(df, column_name, filter_type, filter_value, minimum_acceptable_size=0):
if filter_type == '==':
df_filtered = df[df[column_name]==filter_value]
elif filter_type == '>=':
df_filtered = df[df[column_name]>=filter_value]
elif filter_type == '<=':
df_filtered = df[df[column_name]<=filter_value]
elif filter_type == 'contains':
df_filtered = df[df['target'].str.contains(filter_value)]
if df_filtered.size >= minimum_acceptable_size:
return df_filtered
else:
return df
#the first module becomes text1, the second module file1
def greet(size, target, stage, query):
def raised_zero(x):
if x == 0:
return 'Undisclosed'
else:
return x
df_knn = search(df, query)
df_knn['raised'] = df_knn['raised'].apply(lambda x : raised_zero(x))
df_size = filter_df(df_knn, 'size', '==', size, 1)
if stage != 'ALL':
df_stage = filter_df(df_size, 'stage', '==', stage.lower(), 1)
else:
#we bypass the filter
df_stage = df_size
print(df_stage.size)
df_target = filter_df(df_stage, 'target', 'contains', target, 1)
# display(df_stage)
# df_raised = df_target[(df_target['raised'] >= raised) | (df_target['raised'] == 0)]
#we live the sorting for last
return df_target[0:100] #.sort_values('raised', ascending=False)
with gr.Blocks(theme=gr.themes.Soft(primary_hue='amber', secondary_hue='gray', neutral_hue='amber')) as demo:
gr.Markdown(
"""
# Startup Search Engine
"""
)
size = gr.Radio(['1-10', '11-50', '51-200', '201-500', '500+', '11-500+'], multiselect=False, value='11-500+', label='size')
target = gr.Radio(['B2B', 'B2C', 'B2G', 'B2B2C'], multiselect=False, value='B2B', label='target')
stage = gr.Radio(['pre-seed', 'A', 'B', 'C', 'ALL'], multiselect=False, value='ALL', label='stage')
# raised = gr.Slider(0, 20, value=5, step_size=1, label="Minimum raising (in Millions)")
query = gr.Textbox(label='Describe the Startup you are searching for', value='age reversing')
btn = gr.Button(value="Search for a Startup")
output1 = gr.DataFrame(label='value')
# btn.click(greet, inputs='text', outputs=['dataframe'])
btn.click(greet, [size, target, stage, query], [output1])
demo.launch(share=False)