i cant use git for the life of me. might need more testing

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,12 @@
----
-title: RAG On Images
-emoji: 🔥
-colorFrom: yellow
-colorTo: red
-sdk: gradio
-sdk_version: 4.37.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Image RAG
+emoji: 🔥
+colorFrom: yellow
+colorTo: red
+sdk: gradio
+sdk_version: 4.37.1
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,47 +1,56 @@
-import gradio as gr
-import spaces
-import torch
-from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
-from datasets import load_dataset
-
-dataset = load_dataset("not-lain/embedded-pokemon", split="train")
-dataset = dataset.add_faiss_index("embeddings")
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14")
-model = AutoModelForZeroShotImageClassification.from_pretrained(
-    "openai/clip-vit-large-patch14", device_map=device
-)
-
-
-@spaces.GPU
-def search(query: str, k: int = 4):
-    """a function that embeds a new image and returns the most probable results"""
-
-    pixel_values = processor(images=query, return_tensors="pt")[
-        "pixel_values"
-    ]  # embed new image
-    pixel_values = pixel_values.to(device)
-    img_emb = model.get_image_features(pixel_values)[0]  # because 1 element
-    img_emb = img_emb.cpu().detach().numpy()  # because datasets only works with numpy
-
-    scores, retrieved_examples = dataset.get_nearest_examples(  # retrieve results
-        "embeddings",
-        img_emb,  # compare our new embedded query with the dataset embeddings
-        k=k,  # get only top k results
-    )
-    images = retrieved_examples["image"]
-    # labels = {}
-    # for i in range(k):
-    #     labels[retrieved_examples["text"][k-i]] = scores[k-i]
-
-    return images #, labels
-
-demo = gr.Interface(search, inputs="image", outputs=["gallery"
-                                                     #, "label"
-                                                     ],
-                                                     examples=["./charmander.jpg"],
-)
-
-demo.launch(debug=True)
+import gradio as gr
+from utils import dataset_rag
+
+dirty_hack = True
+
+if dirty_hack:
+    import os    
+    os.environ['KMP_DUPLICATE_LIB_OK']='True'
+
+
+datasets = [
+    "not-lain/embedded-pokemon"
+]
+
+space_installed = None
+
+try:
+    import spaces
+    space_installed = True
+except ImportError:
+    space_installed = False
+
+if space_installed:
+    @spaces.GPU
+    def instance(dataset_name):
+        return dataset_rag.Instance(dataset_name)
+else:
+    def instance(dataset_name):
+        return dataset_rag.Instance(dataset_name)
+
+def download(dataset):
+    global ds
+    client = instance(datasets[0])
+    ds = client
+    return client
+
+def search_ds(image):
+    scores, retrieved_examples = ds.search(image)
+    return retrieved_examples, scores
+
+with gr.Blocks(title="Image RAG") as demo:
+    ds = None
+    interactive_mode = False
+    dataset_name = gr.Dropdown(label="Dataset", choices=datasets, value=datasets[0])
+    download_dataset = gr.Button("Download Dataset")
+
+    search = gr.Image(label="Search Image")
+    search_button = gr.Button("Search")
+    results = gr.Gallery(label="Results")
+    scores = gr.Textbox(label="Scores", type="text", value="")
+    search_button.click(search_ds, inputs=[search], outputs=[results, scores])
+    
+    download_dataset.click(download, dataset_name)
+
+
+demo.launch()	

requirements.txt

@@ -1,7 +1,8 @@
-pillow
-datasets 
-torch
-spaces
-accelerate 
-faiss-cpu
-transformers
+datasets
+accelerate
+loadimg
+faiss-cpu
+numpy==1.26.0
+transformers # hf spaces already have it installed.
+pillow
+gradio # duh

utils/dataset_rag.py

@@ -0,0 +1,64 @@
+from datasets import load_dataset
+import torch
+from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
+from loadimg import load_img
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu' # we should rlly check for mps but, who uses macs (this is a space. lol)
+
+processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14")
+model = AutoModelForZeroShotImageClassification.from_pretrained("openai/clip-vit-large-patch14", device_map = device)
+
+class Instance:
+    def __init__(self, dataset, token=None, split="train"):
+        self.dataset = dataset
+        self.token = token
+        self.split = split
+        self.data = load_dataset(self.dataset, split=self.split)
+        self.data = self.data.add_faiss_index("embeddings")
+    
+    def embed(batch):
+        """a function that embeds a batch of images and returns the embeddings intended for embedding already existing images in an external dataset. (unused)"""
+        pixel_values = processor(images = batch["image"], return_tensors="pt")['pixel_values']
+        pixel_values = pixel_values.to(device)
+        img_emb = model.get_image_features(pixel_values)
+        batch["embeddings"] = img_emb
+        return batch
+    
+    def search(self, query: str, k: int = 3 ):
+        """
+        A function that embeds a query image and returns the most probable results.
+        
+        Args:
+            query: the image to search for
+            k: the number of results to return
+        
+        Returns:
+            scores: the scores of the retrieved examples (cosine similarity i think in this case)
+            retrieved_examples: the retrieved examples
+        """
+
+        pixel_values = processor(images = query, return_tensors="pt")['pixel_values']
+        pixel_values = pixel_values.to(device)
+        img_emb = model.get_image_features(pixel_values)[0]
+        img_emb = img_emb.cpu().detach().numpy()
+
+        scores, retrieved_examples = self.data.get_nearest_examples(
+            "embeddings", img_emb,
+            k=k
+        )
+
+        return scores, retrieved_examples
+
+    def high_level_search(self, img):
+        """
+        High level wrapper for the search function.
+
+        Args:
+            img: input image (path, url, pillow or numpy)
+        
+        Returns:
+            scores: the scores of the retrieved examples (cosine similarity i think in this case)
+            retrieved_examples: the retrieved examples
+        """
+        image = load_img(img)
+        scores, retrieved_examples = self.search(image)