# Initialize a retriever using Qdrant and SentenceTransformer embeddings
from langchain_community.vectorstores import Qdrant
from langchain_community.embeddings import SentenceTransformerEmbeddings
from qdrant_client import QdrantClient
from transformers import AutoModel, AutoTokenizer
import os
import pandas as pd
import gradio as gd

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=vector_db_key,
) # Persists changes to disk, fast prototyping

COLLECTION_NAME="semantic_image_search"


dense_vector_retriever = Qdrant(client, COLLECTION_NAME, embeddings)
images_data = pd.read_csv("/kaggle/input/fashion-product-images-dataset/fashion-dataset/images.csv", on_bad_lines='skip')

def get_link(query):
    Search_Query = query
    neutral_retiever = 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()