Fashion_Vectron / app.py
Cyanido's picture
Update app.py
1ee0865 verified
raw
history blame
1.61 kB
# Initialize a retriever using Qdrant and SentenceTransformer embeddings
from langchain_community.vectorstores import Qdrant
from langchain_community.retrievers.qdrant_sparse_vector_retriever import QdrantSparseVectorRetriever
from langchain_community.embeddings import SentenceTransformerEmbeddings
from qdrant_client import QdrantClient
import pandas as pd
import gradio as gr
embeddings = SentenceTransformerEmbeddings(model_name='sentence-transformers/clip-ViT-B-32')
def get_results(search_results):
filtered_img_ids = [doc.metadata.get("image_id") for doc in search_results]
return filtered_img_ids
client = QdrantClient(
url="https://763bc1da-0673-4535-91ac-b5538ec0287f.us-east4-0.gcp.cloud.qdrant.io:6333",
api_key='UOqiBgqhhu8BBWP98mwjGl7h4IhL2vMAqzO4EI9PEB66A50n9GoIiQ',
) # Persists changes to disk, fast prototyping
COLLECTION_NAME="semantic_image_search"
dense_vector_retriever = Qdrant(client, COLLECTION_NAME, embeddings)
images_data = pd.read_csv("images.csv", on_bad_lines='skip')
def get_link(query):
Search_Query = query
neutral_retiever = QdrantSparseVectorRetriever(retrievers=[dense_vector_retriever.as_retriever()])
result = neutral_retiever.get_relevant_documents(Search_Query)
filtered_images = get_results(result)
filtered_img_ids = [doc.metadata.get("image_id") for doc in result]
links = [images_data.loc[id, 'link'] for id in filtered_img_ids]
# final = '[' + ','.join(links) + ']'
return links
# print(get_link("black shirt for men"))
gr.Interface(fn = get_link, inputs = 'textbox', outputs = 'textbox').launch()