Sync from GitHub Actions

config/model.yaml

@@ -1,10 +1,8 @@
 model:
   name: "openai/clip-vit-large-patch14"
-  learning_rate: 1e-5
-  batch_size: 128
-  epochs: 5
-  grad_accum_steps: 1
-  num_warmup_steps : 100
+  learning_rate: 1e-4
+  batch_size: 32
+  epochs: 10
   new_model: "mohsin416/clip-vit-large-patch14-fashion-retrieval-lora"
 
 index:

config/schema.yaml

@@ -3,8 +3,6 @@ model:
   learning_rate : float
   batch_size : int
   epochs : int
-  num_warmup_steps : int
-  grad_accum_steps : int
   lora_model: str
 
 

visual_product_search/data/dataset.py

@@ -1,23 +1,41 @@
+import os
+import sys
 import torch
-from torch.utils.data import Dataset, DataLoader
-from torch.nn.functional import pad
-from transformers import CLIPProcessor
+import pandas as pd
+from torch.utils.data import Dataset
 from PIL import Image
-from visual_product_search.logger import logging
+from transformers import CLIPProcessor
 from visual_product_search.exception import ExceptionHandle
-import pandas as pd
-import sys, os
+from visual_product_search.logger import logging
+
 
 class ProductDataset(Dataset):
-    def __init__(self, df: pd.DataFrame, processor: CLIPProcessor, img_dir: str, preprocessed_dir: str = None):
-        self.df = df
-        self.processor = processor
-        self.img_dir = img_dir
-        self.preprocessed_dir = preprocessed_dir
+    def __init__(self, df: pd.DataFrame, image_folder: str, processor: CLIPProcessor, preprocessed_dir: str = None):
+        try:
+            df["image_path"] = df["filename"].astype(str).apply(
+                lambda x: os.path.join(image_folder, x)
+            )
+            df = df[df["image_path"].apply(os.path.exists)].reset_index(drop=True)
 
-        if preprocessed_dir:
-            os.makedirs(preprocessed_dir, exist_ok=True)
-            logging.info(f"Preprocessed images will be cached in {preprocessed_dir}")
+            df["caption"] = (
+                df["gender"].fillna("") + " "
+                + df["masterCategory"].fillna("") + " "
+                + df["subCategory"].fillna("") + " "
+                + df["baseColour"].fillna("") + " "
+                + df["articleType"].fillna("") + " "
+                + df["productDisplayName"].fillna("")
+            ).str.strip()
+
+            self.df = df.reset_index(drop=True)
+            self.processor = processor
+            self.preprocessed_dir = preprocessed_dir
+
+            if preprocessed_dir:
+                os.makedirs(preprocessed_dir, exist_ok=True)
+                logging.info(f"Preprocessed images will be cached in {preprocessed_dir}")
+
+        except Exception as e:
+            raise ExceptionHandle(e, sys)
 
     def __len__(self):
         return len(self.df)
@@ -25,44 +43,50 @@ class ProductDataset(Dataset):
     def __getitem__(self, idx):
         try:
             row = self.df.iloc[idx]
-            img_name = row['filename']
-            img_path = os.path.join(self.img_dir, img_name)
+            img_path = row["image_path"]
+
+            tensor_path = (
+                os.path.join(self.preprocessed_dir, row["filename"] + ".pt")
+                if self.preprocessed_dir
+                else None
+            )
 
-            tensor_path = os.path.join(self.preprocessed_dir, img_name + ".pt") if self.preprocessed_dir else None
-                
             if tensor_path and os.path.exists(tensor_path):
                 pixel_values = torch.load(tensor_path)
-                    
             else:
                 try:
                     image = Image.open(img_path).resize((224, 224))
-                    pixel_values = self.processor(images=image, return_tensors="pt")['pixel_values'].squeeze(0)
+                    pixel_values = self.processor(images=image, return_tensors="pt")[
+                        "pixel_values"
+                    ].squeeze(0)
                     if tensor_path:
                         torch.save(pixel_values, tensor_path)
-                        
                 except FileNotFoundError:
                     logging.warning(f"Image not found: {img_path}, using dummy tensor.")
                     pixel_values = torch.zeros(3, 224, 224)
 
-            filtered_row = row.drop(labels=["filename", "link"])
-            text = " ".join(str(v) for v in filtered_row.values)
+            caption = row["caption"]
+            img_link = row.get("link", None)
+
             text_inputs = self.processor.tokenizer(
-                text,
+                caption,
                 padding="max_length",
                 truncation=True,
                 max_length=77,
-                return_tensors="pt"
+                return_tensors="pt",
             )
-            
+
             input_ids = text_inputs["input_ids"].squeeze(0)
             attention_mask = text_inputs["attention_mask"].squeeze(0)
 
             return {
                 "input_ids": input_ids,
                 "attention_mask": attention_mask,
-                "pixel_values": pixel_values
+                "pixel_values": pixel_values,
+                "caption": caption,
+                "img_link": img_link,
             }
 
         except Exception as e:
-            logging.error(f"Error processing item {idx}")
+            logging.error(f"Error processing dataset item {idx}")
             raise ExceptionHandle(e, sys)

visual_product_search/embeddings/train.py

@@ -1,126 +1,80 @@
 import torch
-from torch.nn import CrossEntropyLoss
+from torch import nn
 from torch.amp import GradScaler, autocast
 from torch.optim import AdamW
-from transformers import get_cosine_schedule_with_warmup
+from peft import LoraConfig, get_peft_model
 from visual_product_search.logger import logging
+import time
 from visual_product_search.exception import ExceptionHandle
 import sys
 
-def train(model, dataloader, device, epochs=5, lr=1e-5, grad_accum_steps=1, num_warmup_step=50):
-    optimizer = AdamW(model.parameters(), lr=lr)
-    scaler = GradScaler(device="cuda")
-    total_steps = epochs * len(dataloader)
-    
-    scheduler = get_cosine_schedule_with_warmup(
-        optimizer,
-        num_warmup_steps=num_warmup_step,
-        num_training_steps=total_steps
-    )
-    
+
+def train(model, dataloader, device, epochs=10, lr=1e-4):
     try:
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=16,
+            target_modules=["q_proj", "v_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="FEATURE_EXTRACTION"
+        )
+        model = get_peft_model(model, lora_config)
+        model.to(device)
+
+        optimizer = AdamW(model.parameters(), lr=lr)
+        loss_fn = nn.CrossEntropyLoss()
+        scaler = GradScaler()
+
         for epoch in range(epochs):
             model.train()
-            total_loss = 0
-            
-            for step, batch in enumerate(dataloader):
-                try:
-                    input_ids = batch['input_ids'].to(device)
-                    attention_mask = batch['attention_mask'].to(device)
-                    pixel_values = batch['pixel_values'].to(device)
-                    
-                    optimizer.zero_grad()
-                    
-                    with autocast(device_type="cuda"):
-                        outputs = model(input_ids=input_ids,
-                                        attention_mask=attention_mask,
-                                        pixel_values=pixel_values)
-                        
-                        img_embd = outputs.image_embeds
-                        text_embd = outputs.text_embeds
-                        
-                        img_embd = img_embd / img_embd.norm(p=2, dim=-1, keepdim=True)
-                        text_embd = text_embd / text_embd.norm(p=2, dim=-1, keepdim=True)
-                        
-                        logits_per_image = img_embd @ text_embd.T
-                        labels = torch.arange(len(img_embd)).to(device)
+            total_loss = 0.0
+            start_epoch = time.time()
+            logging.info(f"--------- Epoch {epoch+1}/{epochs} ---------")
 
-                        loss_i = CrossEntropyLoss()(logits_per_image, labels)
-                        loss_t = CrossEntropyLoss()(logits_per_image.T, labels)
-                        
-                        loss = (loss_i + loss_t) / 2
-                    
-                    loss = loss / grad_accum_steps
-                    scaler.scale(loss).backward()
-                    
-                    if (step + 1) % grad_accum_steps == 0:
-                        scaler.step(optimizer)
-                        scaler.update()
-                        optimizer.zero_grad()
-                        scheduler.step()
-                        
-                    total_loss += loss.item() * grad_accum_steps 
-                    
-                    if step % 10 == 0:
-                        logging.info(f"Epoch {epoch+1}/{epochs}, Step {step}, Loss: {loss.item():.4f}")
-                
-                except Exception as e:
-                    logging.error(f"Failure at step {step} in epoch {epoch+1}")
-                    raise ExceptionHandle(e, sys)
-            
-            
-            avg_loss = total_loss / len(dataloader)
-            logging.info(f" Epoch {epoch + 1} / {epochs} finished | Avg Loss : {avg_loss:.4f}")
-        
-        model.eval()
-        return model
-            
-    except Exception as e:
-        logging.critical("Training loop crashed")
-        raise ExceptionHandle(e, sys)
-    
-    
-    
-    
-# for epoch in range(EPOCHS):
-#     model.train()
-#     total_loss = 0.0
-#     start_epoch = time.time()
-#     print(f"\n======== Epoch {epoch+1}/{EPOCHS} ========")
-    
-#     for batch_idx, (imgs, captions) in enumerate(dataloader, start=1):
-#         batch_start = time.time()
-#         imgs = imgs.to(DEVICE)
-#         inputs_txt = processor(text=list(captions), return_tensors="pt", padding=True, truncation=True).to(DEVICE)
+            for batch_idx, batch in enumerate(dataloader, start=1):
+                imgs = batch["pixel_values"].to(device)
+                input_ids = batch["input_ids"].to(device)
+                attention_mask = batch["attention_mask"].to(device)
 
-#         optimizer.zero_grad()
-#         with autocast(device_type="cuda", dtype=torch.float16):
-#             inputs_img = processor(images=imgs, return_tensors="pt").to(DEVICE)
-#             img_embeds = model.get_image_features(**inputs_img)
-#             txt_embeds = model.get_text_features(**inputs_txt)
+                optimizer.zero_grad()
+                with autocast(device_type="cuda", dtype=torch.float16):
+                    img_embeds = model.get_image_features(pixel_values=imgs)
+                    txt_embeds = model.get_text_features(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask
+                    )
 
-#             img_embeds = nn.functional.normalize(img_embeds, dim=-1)
-#             txt_embeds = nn.functional.normalize(txt_embeds, dim=-1)
+                    img_embeds = nn.functional.normalize(img_embeds, dim=-1)
+                    txt_embeds = nn.functional.normalize(txt_embeds, dim=-1)
 
-#             logits = img_embeds @ txt_embeds.T * 100
-#             labels = torch.arange(len(logits), device=DEVICE)
-#             loss_i2t = loss_fn(logits, labels)
-#             loss_t2i = loss_fn(logits.T, labels)
-#             loss = (loss_i2t + loss_t2i) / 2
+                    logits = img_embeds @ txt_embeds.T * 100
+                    labels = torch.arange(len(logits), device=device)
+                    loss_i2t = loss_fn(logits, labels)
+                    loss_t2i = loss_fn(logits.T, labels)
+                    loss = (loss_i2t + loss_t2i) / 2
 
-#         scaler.scale(loss).backward()
-#         scaler.step(optimizer)
-#         scaler.update()
+                scaler.scale(loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
 
-#         total_loss += loss.item()
+                total_loss += loss.item()
 
-#         if batch_idx % 50 == 0 or batch_idx == len(dataloader):
-#             elapsed = time.time() - batch_start
-#             print(f"[Epoch {epoch+1} Batch {batch_idx}/{len(dataloader)}] "
-#                   f"Batch Loss: {loss.item():.4f}, "
-#                   f"Elapsed: {elapsed:.2f}s, "
-#                   f"GPU Memory Used: {torch.cuda.memory_allocated()/1024**3:.2f} GB")
+                if batch_idx % 50 == 0 or batch_idx == len(dataloader):
+                    logging.info(
+                        f"[Epoch {epoch+1} Batch {batch_idx}/{len(dataloader)}] "
+                        f"Batch Loss: {loss.item():.4f}"
+                    )
+
+            avg_loss = total_loss / len(dataloader)
+            epoch_time = time.time() - start_epoch
+            logging.info(
+                f"Epoch {epoch+1} Completed | Avg Loss: {avg_loss:.4f} | "
+                f"Time: {epoch_time/60:.2f} min"
+            )
+            
+        return model
 
-#     avg_loss = total_loss / len(dataloader)
-#     epoch_time = time.time() - start_epoch
-#     print(f"Epoch {epoch+1} Completed | Avg Loss: {avg_loss:.4f} | Time: {epoch_time/60:.2f} min")
+    except Exception as e:
+        logging.critical("Training loop crashed")
+        raise ExceptionHandle(e, sys)

visual_product_search/indexing/indexer.py

@@ -47,7 +47,7 @@ class DatabaseIndexer:
                 index_params = {
                     "index_type" : "HNSW",
                     "metric_type" : "COSINE",
-                    "params" : {"M" : 8, "efConstruction" : 64}
+                    "params" : {"M" : 48, "efConstruction" : 200}
                 }
                 self.collection.create_index(field_name="embedding", index_params=index_params)
                 logging.info("Index Created successfully")

visual_product_search/pipeline/training_pipeline.py

@@ -1,6 +1,6 @@
 import torch, gc
 from torch.utils.data import DataLoader
-import numpy as np
+import torch.nn.functional.normalize as F
 import os
 from pathlib import Path
 import sys
@@ -58,8 +58,6 @@ class VisualProductPipeline:
                 device,
                 epochs=self.config["model"]["epochs"],
                 lr=self.config["model"]["learning_rate"],
-                grad_accum_steps=self.config["model"]["grad_accum_steps"],
-                num_warmup_step=self.config["model"]["num_warmup_steps"]
             )
             logging.info("Model training completed")
             return trained_model
@@ -67,26 +65,43 @@ class VisualProductPipeline:
         except Exception as e:
             raise ExceptionHandle(e, sys)
     
-    def create_embeddings(self, df, img_dir, model, processor, device):
+    def create_embeddings(self, dataloader, model, device):
         try:
             logging.info("Creating embeddings")
             embeddings = []
-            metadata = []
             img_link = []
+            metadata = []
             
-            for row in df.itertuples(index=False):
-                img_path = f"{img_dir}/{row.filename}"
-                embd =  get_image_embedding(img_path, model, processor, device)
-                embeddings.append(embd)
-                img_link.append(row.link)
+            with torch.no_grad():
+                for batch in dataloader:
+                    try:
+                        imgs = batch["pixel_values"].to(device)
+                        caps = batch["caption"]
+                        links = batch["img_link"]
+
+                        if imgs is None or len(imgs) == 0:
+                            print("Image not found, skipping")
+                            continue
+
+                        img_embeds = model.get_image_features(pixel_values=imgs)
+                        img_embeds = F(img_embeds, dim=-1)
+
+                        embeddings.append(img_embeds.cpu())
+                        metadata.extend(caps)
+                        img_link.extend(links)
+
+                    except Exception as e:
+                        print(f"Skipping batch due to error: {e}")
+                        continue
+
+            if embeddings:
+                embeddings = torch.cat(embeddings, dim=0)
+            else:
+                embeddings = torch.empty(0)
                 
-                filtered_row = row._asdict()
-                filtered_row.pop("filename", None)
-                filtered_row.pop("link", None)
-                metadata.append(" ".join(str(v) for v in filtered_row.values()))
-            
-            embeddings = np.vstack(embeddings)
-            logging.info(f"create embeddings for {len(df)} samples")
+            logging.info(f"Embeddings shape: {embeddings.shape}")
+            logging.info(f"Metadata length: {len(metadata)}")
+            logging.info(f"Image links length: {len(img_link)}")
             return embeddings, metadata, img_link
         
         except Exception as e:
@@ -133,7 +148,7 @@ class VisualProductPipeline:
             df, img_dir = self.data_ingestion()
             model, processor, device = self.model_loading()
             
-            dataset = ProductDataset(df, processor, img_dir, str(self.cache_dir))
+            dataset = ProductDataset(df, img_dir, processor, str(self.cache_dir))
             
             dataloader = DataLoader(
                 dataset,
@@ -148,7 +163,7 @@ class VisualProductPipeline:
             trained_model = self.start_training(model, dataloader, device)
             self.push_hub(trained_model, processor)
             
-            embeddings, metadata, img_link = self.create_embeddings(df, img_dir, trained_model, processor, device)
+            embeddings, metadata, img_link = self.create_embeddings(dataloader, trained_model, device)
             self.start_indexing(embeddings, metadata, img_link)
             
             del model, processor, trained_model