Sarat Chandra
initial commit
ec14c0e
raw
history blame contribute delete
No virus
1.32 kB
from flask import Flask, render_template
from flask_socketio import SocketIO
import pandas as pd
import numpy as np
import google.generativeai as genai
# import KEYS
import os
app = Flask(__name__)
socketio = SocketIO(app)
# genai.configure(api_key=KEYS.api_key.GOOGLE_API_KEY)
genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
n_steps = 10
df = pd.DataFrame()
for i in range(0,n_steps):
df = pd.concat([df,pd.read_feather(f"data_{i}.feather")])
df.reset_index(inplace=True,drop=True)
def get_results(top_n = 6,query = "men shirt"):
query_embedding = genai.embed_content(model="models/text-embedding-004",
content=query,
task_type="retrieval_query")['embedding']
scores = df['embedding'].apply(lambda x: np.dot(x,query_embedding))
scores = scores.sort_values(ascending=False)[0:top_n]
return df.loc[scores.index][['productDisplayName','link']].to_numpy()
@app.route('/')
def index():
return render_template('index.html')
@socketio.on("search")
def get_products(query):
data = []
for x in get_results(query=query):
data.append({'url':x[1],'name':x[0]})
socketio.emit('data',data)
if __name__ == '__main__':
socketio.run(app,port=7860,allow_unsafe_werkzeug=True,host='0.0.0.0')