feat: feluda store supports audio (#78)

* feat: feluda store supports audio

* fix: delete and refresh for ES

* dhore: profiling audio operator

.gitignore

@@ -135,4 +135,4 @@ src/api/core/operators/yolov8n-seg.pt
 # Audio Files
 **100_audio_files/
 **50_audio_files/
-
+**audio_files/

src/api/config-server.yml

@@ -6,6 +6,7 @@ store :
     image_index_name : "image"
     text_index_name : "text"
     video_index_name : "video"
+    audio_index_name : "audio"
 
 queue : 
   label : "Queue"

src/api/core/config.py

@@ -22,6 +22,7 @@ class StoreParameters:
     image_index_name: str
     text_index_name: str
     video_index_name: str
+    audio_index_name: str
 
 
 @dataclass

src/api/core/models/media.py

@@ -13,6 +13,7 @@ class MediaType(Enum):
     TEXT = "text"
     IMAGE = "image"
     VIDEO = "video"
+    AUDIO = "audio"
 
     @classmethod
     def make(media_type):
@@ -22,5 +23,7 @@ def make(media_type):
             return MediaType.IMAGE
         elif media_type is "video":
             return MediaType.VIDEO
+        elif media_type is "audio":
+            return MediaType.AUDIO
         else:
             return MediaType.UNSUPPORTED

src/api/core/operators/audio_vec_embedding_requirements.in

@@ -1,3 +1,5 @@
 numpy==1.26.4
 librosa==0.10.1
-panns-inference==0.1.1
+panns-inference==0.1.1
+torch==2.2.0+cpu
+torchvision==0.17.0+cpu

src/api/core/operators/audio_vec_embedding_requirements.txt

@@ -2,8 +2,10 @@
 # This file is autogenerated by pip-compile with Python 3.12
 # by the following command:
 #
-#    pip-compile audio_vec_embedding_requirements.in
+#    pip-compile --find-links=https://download.pytorch.org/whl/torch_stable.html audio_vec_embedding_requirements.in
 #
+--find-links https://download.pytorch.org/whl/torch_stable.html
+
 audioread==3.0.1
     # via librosa
 certifi==2024.2.2
@@ -18,10 +20,16 @@ cycler==0.12.1
     # via matplotlib
 decorator==5.1.1
     # via librosa
+filelock==3.13.1
+    # via torch
 fonttools==4.48.1
     # via matplotlib
+fsspec==2024.2.0
+    # via torch
 idna==3.6
     # via requests
+jinja2==3.1.3
+    # via torch
 joblib==1.3.2
     # via
     #   librosa
@@ -37,10 +45,16 @@ librosa==0.10.1
     #   torchlibrosa
 llvmlite==0.42.0
     # via numba
+markupsafe==2.1.5
+    # via jinja2
 matplotlib==3.8.2
     # via panns-inference
+mpmath==1.3.0
+    # via sympy
 msgpack==1.0.7
     # via librosa
+networkx==3.2.1
+    # via torch
 numba==0.59.0
     # via librosa
 numpy==1.26.4
@@ -54,14 +68,17 @@ numpy==1.26.4
     #   scipy
     #   soxr
     #   torchlibrosa
+    #   torchvision
 packaging==23.2
     # via
     #   matplotlib
     #   pooch
 panns-inference==0.1.1
     # via -r audio_vec_embedding_requirements.in
 pillow==10.2.0
-    # via matplotlib
+    # via
+    #   matplotlib
+    #   torchvision
 platformdirs==4.2.0
     # via pooch
 pooch==1.8.0
@@ -73,7 +90,9 @@ pyparsing==3.1.1
 python-dateutil==2.8.2
     # via matplotlib
 requests==2.31.0
-    # via pooch
+    # via
+    #   pooch
+    #   torchvision
 scikit-learn==1.4.0
     # via librosa
 scipy==1.12.0
@@ -86,11 +105,21 @@ soundfile==0.12.1
     # via librosa
 soxr==0.3.7
     # via librosa
+sympy==1.12
+    # via torch
 threadpoolctl==3.2.0
     # via scikit-learn
+torch==2.2.0+cpu
+    # via
+    #   -r audio_vec_embedding_requirements.in
+    #   torchvision
 torchlibrosa==0.1.0
     # via panns-inference
+torchvision==0.17.0+cpu
+    # via -r audio_vec_embedding_requirements.in
 typing-extensions==4.9.0
-    # via librosa
+    # via
+    #   librosa
+    #   torch
 urllib3==2.2.0
     # via requests

src/api/core/store/es_vec.py

@@ -19,6 +19,7 @@ def __init__(self, config: StoreConfig):
             "text": config.parameters.text_index_name,
             "image": config.parameters.image_index_name,
             "video": config.parameters.video_index_name,
+            "audio": config.parameters.audio_index_name,
         }
 
     def connect(self):
@@ -57,7 +58,7 @@ def optionally_create_index(self):
 
     def delete_indices(self):
         for index in self.indices:
-            self.client.indices.delete(index=index)
+            self.client.indices.delete(index=self.indices[index])
 
     def get_indices(self):
         index_list = ""
@@ -78,7 +79,7 @@ def store(self, media_type: MediaType, doc):
 
     def refresh(self):
         for index in self.indices:
-            self.client.indices.refresh(index=index)
+            self.client.indices.refresh()
 
     def find(self, index_name, vec):
         if type(vec) == np.ndarray:
@@ -91,6 +92,8 @@ def find(self, index_name, vec):
             calculation = "1 / (1 + l2norm(params.query_vector, 'image_vec'))"
         elif index_name == self.indices["video"]:
             calculation = "1 / (1 + l2norm(params.query_vector, 'vid_vec'))"
+        elif index_name == self.indices["audio"]:
+            calculation = "1 / (1 + l2norm(params.query_vector, 'audio_vec'))"
 
         print("calculation:", calculation)
         q = {

src/api/core/store/es_vec_mappings.py

@@ -87,4 +87,26 @@
                     }
                 }
             }""",
+    "audio": """{
+                "mappings": {
+                    "properties":{
+                        "e_kosh_id": {
+                            "type": "keyword"
+                        },
+                        "dataset": {
+                            "type": "keyword"
+                        },
+                        "metadata": {
+                            "type": "object"
+                        },
+                        "audio_vec": {
+                            "type":"dense_vector",
+                            "dims": 2048
+                        },
+                        "date_added": {
+                            "type": "date"
+                        }
+                    }
+                }
+            }""",    
 }