import numpy as np
from PIL import Image
from araclip import AraClip
import gradio as gr
model = AraClip.from_pretrained("Arabic-Clip/araclip")

def search(labels, image):
    # process labels
    labels = labels.split(",")
    labels =  [item.strip() for item in labels if item != ""]
    
    # embed data 
    image_features = model.embed(image=image)
    text_features = np.stack([model.embed(text=label) for label in labels])
    # search for most similar data
    similarities = text_features @ image_features
    best_match = labels[np.argmax(similarities)]
    return best_match

    

# data
labels = ["قطة جالسة", "قطة تقفز" ,"كلب", "حصان"]
image = Image.open("cat.png")

demo = gr.Interface(search,["text","image"],["text"],examples=[["حصان, كلب, قطة", "cat.png"]])

demo.launch(debug=True)