Spaces:
Running
Running
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() | |
def index(): | |
return render_template('index.html') | |
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') |