initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+data/embeddings.bin filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+.env
+*.pyc
+*.py[co]
+encoder.bin
+.DS_Store

app.py

@@ -0,0 +1,48 @@
+import streamlit as st
+import random
+import os
+import dotenv
+from datasets import load_dataset, concatenate_datasets
+from recommendation import get_recommendations
+
+dotenv.load_dotenv()
+
+START = random.randint(a=1,b=200)
+END = START + 10
+
+dataset = load_dataset("Ransaka/youtube_recommendation_data", token=os.environ.get('HF'))
+dataset = concatenate_datasets([dataset['train'], dataset['test']])
+
+
+
+pil_images = dataset['image'][START:END]
+
+def show_image_metadata_and_related_info(image_index):
+    selected_image = pil_images[image_index]
+    image_title = dataset['title'][image_index]
+    st.image(selected_image, caption=f"{image_title}", use_column_width=True)
+    
+    with st.expander("You May Also Like.."):
+        dataset_s = get_recommendations(selected_image,image_title)
+        
+        col1_row1, col2_row1 = st.columns(2)
+        with col1_row1:
+            st.image(image=dataset_s['image'][0], caption=dataset_s['title'][0], width=200)
+        with col2_row1:
+            st.image(image=dataset_s['image'][1], caption=dataset_s['title'][1], width=200)
+
+        # Second Row
+        col1_row2, col2_row2 = st.columns(2)
+        with col1_row2:
+            st.image(image=dataset_s['image'][2], caption=dataset_s['title'][2], width=200)
+        with col2_row2:
+            st.image(image=dataset_s['image'][3], caption=dataset_s['title'][3], width=200)
+
+def main():
+    st.title("Youtube Recommendation Engine")
+
+    for i, image in enumerate(pil_images):
+        show_image_metadata_and_related_info(i)
+
+if __name__ == '__main__':
+    main()

indexer.py

@@ -0,0 +1,39 @@
+import faiss
+
+class Indexer:
+  def __init__(self, embed_vec):
+    self.embeddings_vec = embed_vec
+    self.build_index()
+
+  def build_index(self):
+    """
+    Build the index for the embeddings.
+
+    This function initializes the index for the embeddings. It calculates the dimension (self.d)
+    of the embeddings vector and creates an IndexFlatL2 object (self.index) for the given dimension.
+    It then adds the embeddings vector (self.embeddings_vec) to the index.
+
+    Parameters:
+    - None
+
+    Return:
+    - None
+    """
+    self.d = self.embeddings_vec.shape[1]
+    self.index = faiss.IndexFlatL2(self.d)
+    self.index.add(self.embeddings_vec)
+
+  def topk(self, vector, k = 4):
+    """
+        A function that takes in a vector and an optional parameter k and returns the indices of the k nearest neighbors in the index.
+
+        Parameters:
+            vector: A numpy array representing the input vector.
+            k (optional): An integer representing the number of nearest neighbors to retrieve. Defaults to 4 if not specified.
+
+        Returns:
+            I: A numpy array containing the indices of the k nearest neighbors in the index.
+    """
+    # vec = self.retreaver.encode(text)['embeddings'].detach().cpu().numpy()
+    _, I = self.index.search(vector, k)
+    return I

models/encoder.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels=1, out_channels=16, latent_dim=64, act_fn=nn.ReLU()):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1),  # (480, 360)
+            act_fn,
+            nn.Conv2d(out_channels, out_channels, 3, padding=1),
+            act_fn,
+            nn.Conv2d(out_channels, 2 * out_channels, 3, padding=1, stride=2),  # (240, 180)
+            act_fn,
+            nn.Conv2d(2 * out_channels, 2 * out_channels, 3, padding=1),
+            act_fn,
+            nn.Conv2d(2 * out_channels, 4 * out_channels, 3, padding=1, stride=2),  # (120, 90)
+            act_fn,
+            nn.Conv2d(4 * out_channels, 4 * out_channels, 3, padding=1),
+            act_fn,
+            nn.Flatten(),
+            nn.Linear(4 * out_channels * 120 * 90, latent_dim),
+            act_fn
+        )
+
+    def forward(self, x):
+        x = x.view(-1, 1, 480, 360)
+        output = self.net(x)
+        return output

recommendation.py

@@ -0,0 +1,43 @@
+from datasets import load_dataset, concatenate_datasets
+from sentence_transformers import SentenceTransformer
+from torchvision import transforms
+from models.encoder import Encoder
+from indexer import Indexer
+import dotenv
+import torch
+import os
+
+dotenv.load_dotenv()
+
+model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+encoder = Encoder()
+encoder.load_state_dict(torch.load('./models/encoder.bin', map_location=torch.device('cpu')))
+
+dataset = load_dataset("Ransaka/youtube_recommendation_data", token=os.environ.get('HF'))
+dataset = concatenate_datasets([dataset['train'], dataset['test']])
+
+latent_data = torch.load("data/latent_data_final.bin")
+embeddings = torch.load("data/embeddings.bin")
+
+text_embedding_index = Indexer(embeddings)
+image_embedding_index = Indexer(latent_data)
+
+def get_recommendations(image, title):
+#   title = [dataset[product_id]['title']]
+  title_embeds = torch.randn(1,768)#model.encode(title, normalize_embeddings=True)
+  image = transforms.ToTensor()(image.convert("L"))
+  image_embeds =  encoder(image).detach().numpy()
+
+  image_candidates = image_embedding_index.topk(image_embeds)
+  title_candidates = text_embedding_index.topk(title_embeds)
+  final_candidates = []
+  final_candidates.append(list(image_candidates[0]))
+  final_candidates.append(list(title_candidates[0]))
+  final_candidates = sum(final_candidates,[])
+  final_candidates = list(set(final_candidates))
+  results_dict = {"image":[], "title":[]}
+  for candidate in final_candidates:
+    results_dict['image'].append(dataset['image'][candidate])
+    results_dict['title'].append(dataset['title'][candidate])
+  return results_dict

requirements.txt

@@ -0,0 +1,9 @@
+transformers==4.34.1
+sentence_transformers==2.3.1
+faiss-cpu==1.7.4
+pandas==1.5.3
+matplotlib==3.6.0
+streamlit==1.29.0
+torch==2.1.0
+datasets==2.14.5
+torchvision==0.16.0