Sync from GitHub Actions

.gitattributes

@@ -1,35 +0,0 @@
-*.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

.gitignore

@@ -0,0 +1,8 @@
+note.txt
+__pycache__/
+*.pyc
+.ipynb_checkpoints/
+.env
+.venv/
+venv/
+artifacts/

Dockerfile

@@ -1,20 +1,16 @@
-FROM python:3.13.5-slim
+FROM python:3.11-slim
 
 WORKDIR /app
 
-RUN apt-get update && apt-get install -y \
-    build-essential \
-    curl \
-    git \
-    && rm -rf /var/lib/apt/lists/*
+ENV PYTHONUNBUFFERED=1
+ENV PIP_NO_CACHE_DIR=1
 
-COPY requirements.txt ./
-COPY src/ ./src/
+COPY requirements.txt .
 
-RUN pip3 install -r requirements.txt
+RUN pip install --upgrade pip && pip install -r requirements.txt
 
-EXPOSE 8501
+COPY . .
 
-HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
+EXPOSE 8501
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+CMD ["streamlit", "run", "app/app.py", "--server.address=0.0.0.0", "--server.port=8501"]

app/app.py

@@ -0,0 +1,200 @@
+import json
+import sys
+from pathlib import Path
+import streamlit as st
+from PIL import Image
+from inference import load_predictor, render_prediction_card, render_top_predictions, render_metrics
+from autocatalog.utils.config import load_config
+
+ROOT_DIR = Path(__file__).resolve().parents[1]
+if str(ROOT_DIR) not in sys.path:
+    sys.path.insert(0, str(ROOT_DIR))
+    
+
+def main():
+    st.set_page_config(
+        page_title="AutoCatalogAI",
+        page_icon="🛍️",
+        layout="wide",
+    )
+
+    st.markdown(
+        """
+        <style>
+        .main-title {
+            font-size: 2.5rem;
+            font-weight: 800;
+            margin-bottom: 0.2rem;
+        }
+        .subtitle {
+            color: #666;
+            font-size: 1.05rem;
+            margin-bottom: 2rem;
+        }
+        .prediction-card {
+            border: 1px solid #e5e7eb;
+            border-radius: 14px;
+            padding: 16px;
+            margin-bottom: 12px;
+            background: #ffffff;
+            box-shadow: 0 1px 4px rgba(0,0,0,0.04);
+        }
+        .prediction-header {
+            display: flex;
+            justify-content: space-between;
+            align-items: center;
+        }
+        .task-name {
+            font-size: 0.9rem;
+            font-weight: 700;
+            color: #374151;
+        }
+        .confidence {
+            font-size: 0.9rem;
+            font-weight: 700;
+            color: #111827;
+        }
+        .label {
+            font-size: 1.25rem;
+            font-weight: 800;
+            margin-top: 8px;
+            margin-bottom: 10px;
+        }
+        .bar-bg {
+            width: 100%;
+            height: 8px;
+            background: #e5e7eb;
+            border-radius: 999px;
+            overflow: hidden;
+        }
+        .bar-fill {
+            height: 100%;
+            background: #111827;
+            border-radius: 999px;
+        }
+        .catalog-box {
+            border: 1px solid #e5e7eb;
+            border-radius: 14px;
+            padding: 18px;
+            background: #fafafa;
+            margin-bottom: 16px;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    config = load_config("configs/config.yaml")
+    repo_id = config.get("model", {}).get(
+        "repo_id",
+        "mohsin416/autocatalogai-clip-multitask",
+    )
+    
+    top_k = int(config.get("inference", {}).get("top_k", 3))
+    device = config.get("inference", {}).get("device", None)
+    
+    st.markdown('<div class="main-title">AutoCatalogAI</div>', unsafe_allow_html=True)
+    st.markdown(
+        '<div class="subtitle">Fashion product attribute extraction and catalog metadata generation using CLIP multi-task learning.</div>',
+        unsafe_allow_html=True,
+    )
+    
+    with st.sidebar:
+        st.header("Settings")
+        st.write("Model repository")
+        st.code(repo_id)
+        
+        selected_top_k = st.slider(
+            "Top-K predictions",
+            min_value=1,
+            max_value=5,
+            value=top_k,
+        )
+        
+        st.divider()
+        st.caption("Model loads from Hugging Face Hub and runs inference only.")
+    
+    with st.spinner("Loading AutoCatalogAI model..."):
+        predictor = load_predictor(
+            repo_id=repo_id,
+            device=device,
+            top_k=selected_top_k
+        )
+    
+    render_metrics(predictor.get_model_metrics())
+    st.divider()
+    
+    left_col, right_col = st.columns([0.9, 1.1])
+    
+    with left_col:
+        st.subheader("Upload Product Image")
+        uploaded_file = st.file_uploader(
+            "Choose a product image",
+            type=["jpg", "jpeg", "png", "webp"],
+            label_visibility="collapsed",
+        )
+        
+        if uploaded_file is not None:
+            image = Image.open(uploaded_file).convert("RGB")
+            st.image(image, caption="Uploaded Image", use_container_width=True)
+
+    with right_col:
+        st.subheader("Prediction Result")
+
+        if uploaded_file is None:
+            st.info("Upload a fashion product image to generate catalog attributes.")
+            return
+
+        if st.button("Generate Catalog", type="primary", use_container_width=True):
+            with st.spinner("Predicting product attributes..."):
+                result = predictor.predict(
+                    image=image,
+                    top_k=selected_top_k,
+                )
+
+            prediction = result["prediction"]
+            catalog_output = result["catalog_output"]
+            runtime = result["runtime"]
+
+            st.markdown(
+                f"""
+                <div class="catalog-box">
+                    <strong>Suggested Title</strong>
+                    <h3>{catalog_output["suggested_title"]}</h3>
+                </div>
+                """,
+                unsafe_allow_html=True,
+            )
+
+            st.markdown("**Search Tags**")
+            st.write(", ".join(catalog_output["search_tags"]))
+            st.markdown("**Predicted Attributes**")
+
+            for task, task_result in prediction.items():
+                render_prediction_card(task, task_result)
+
+            render_top_predictions(prediction)
+            st.markdown("**Runtime**")
+            st.write(f"Device: `{runtime['device']}`")
+            st.write(f"Inference time: `{runtime['inference_time_ms']:.2f} ms`")
+
+            json_output = json.dumps(
+                catalog_output["json_export"],
+                indent=2,
+                ensure_ascii=False,
+            )
+
+            st.download_button(
+                label="Download JSON",
+                data=json_output,
+                file_name="autocatalogai_prediction.json",
+                mime="application/json",
+                use_container_width=True,
+            )
+
+            with st.expander("Raw JSON Output"):
+                st.json(catalog_output["json_export"])
+
+
+if __name__ == "__main__":
+    main()

app/inference.py

@@ -0,0 +1,62 @@
+import streamlit as st
+from autocatalog.inference.predictor import AutoCatalogPredictor
+
+@st.cache_resource(show_spinner=False)
+def load_predictor(repo_id, device, top_k):
+    return AutoCatalogPredictor(
+        repo_id=repo_id,
+        device=device,
+        top_k=top_k
+    )
+    
+    
+def format_percent(value):
+    return f"{value * 100:.2f}%"
+
+
+def render_prediction_card(task_name, task_result):
+    label = task_result["label"]
+    confidence = task_result["confidence"]
+    
+    st.markdown(
+        f"""
+        <div class="prediction-card">
+            <div class="prediction-header">
+                <span class="task-name">{task_name}</span>
+                <span class="confidence">{format_percent(confidence)}</span>
+            </div>
+            <div class="label">{label}</div>
+            <div class="bar-bg">
+                <div class="bar-fill" style="width: {confidence * 100:.2f}%"></div>
+            </div>
+        </div>
+        """,
+        unsafe_allow_html=True,
+    )
+
+def render_top_predictions(prediction):
+    with st.expander("View Top-3 Predictions"):
+        for task, result in prediction.items():
+            st.markdown(f"**{task}**")
+
+            for item in result["top_3"]:
+                st.write(f"{item['label']} — {format_percent(item['confidence'])}")
+
+            st.divider()
+            
+
+def render_metrics(metrics):
+    if not metrics:
+        return
+
+    overall = metrics.get("overall_metrics", {})
+    if not overall:
+        return
+
+    st.subheader("Model Evaluation")
+    col1, col2, col3, col4 = st.columns(4)
+
+    col1.metric("Average Accuracy", format_percent(overall.get("average_accuracy", 0)))
+    col2.metric("Weighted F1", format_percent(overall.get("average_weighted_f1", 0)))
+    col3.metric("Top-3 Accuracy", format_percent(overall.get("average_top3_accuracy", 0)))
+    col4.metric("Test Samples", f"{overall.get('test_samples', 0):,}")

autocatalog/inference/catalog_generator.py

@@ -0,0 +1,91 @@
+def normalize_text(value):
+    if value is None:
+        return ""
+
+    value = str(value).strip()
+    value = value.replace("_", " ")
+    value = value.replace("-", " ")
+    value = " ".join(value.split())
+
+    return value
+
+
+def generate_title(predicted_labels):
+    parts = []
+    
+    gender = normalize_text(predicted_labels.get("gender"))
+    color = normalize_text(predicted_labels.get("baseColour"))
+    usage = normalize_text(predicted_labels.get("usage"))
+    article_type = normalize_text(predicted_labels.get("articleType"))
+
+    for value in [gender, color, usage, article_type]:
+        if value:
+            parts.append(value)
+
+    return " ".join(parts)
+
+
+def generate_search_tags(predicted_labels):
+    tags = []
+    unique_tags = []
+    
+    gender = normalize_text(predicted_labels.get("gender")).lower()
+    master_category = normalize_text(predicted_labels.get("masterCategory")).lower()
+    sub_category = normalize_text(predicted_labels.get("subCategory")).lower()
+    article_type = normalize_text(predicted_labels.get("articleType")).lower()
+    color = normalize_text(predicted_labels.get("baseColour")).lower()
+    season = normalize_text(predicted_labels.get("season")).lower()
+    usage = normalize_text(predicted_labels.get("usage")).lower()
+
+    if gender and article_type:
+        tags.append(f"{gender} {article_type}")
+
+    if color and article_type:
+        tags.append(f"{color} {article_type}")
+
+    if usage and sub_category:
+        tags.append(f"{usage} {sub_category}")
+
+    if season and master_category:
+        tags.append(f"{season} {master_category}")
+
+    if gender and usage:
+        tags.append(f"{gender} {usage} wear")
+
+    if color and usage:
+        tags.append(f"{color} {usage} fashion")
+
+    if sub_category:
+        tags.append(sub_category)
+
+    if article_type:
+        tags.append(article_type)
+
+    for tag in tags:
+        tag = " ".join(tag.split())
+
+        if tag and tag not in unique_tags:
+            unique_tags.append(tag)
+
+    return unique_tags
+
+
+def generate_catalog_output(predicted_labels):
+    suggested_title = generate_title(predicted_labels)
+    search_tags = generate_search_tags(predicted_labels)
+
+    return {
+        "suggested_title": suggested_title,
+        "search_tags": search_tags,
+        "json_export": {
+            "gender": predicted_labels.get("gender"),
+            "category": predicted_labels.get("masterCategory"),
+            "subcategory": predicted_labels.get("subCategory"),
+            "article_type": predicted_labels.get("articleType"),
+            "color": predicted_labels.get("baseColour"),
+            "season": predicted_labels.get("season"),
+            "usage": predicted_labels.get("usage"),
+            "title": suggested_title,
+            "tags": search_tags,
+        },
+    }

autocatalog/inference/predictor.py

@@ -0,0 +1,198 @@
+import json
+import time
+from pathlib import Path
+
+import torch
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from transformers import CLIPImageProcessor
+
+from autocatalog.models.multitask_clip import CLIPMultiTaskClassifier
+from autocatalog.inference.catalog_generator import generate_catalog_output
+
+
+class AutoCatalogPredictor:
+    def __init__(
+        self,
+        repo_id="mohsin416/autocatalogai-clip-multitask",
+        device=None,
+        top_k=3,
+    ):
+        self.repo_id = repo_id
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.top_k = top_k
+        
+        self.model_path = hf_hub_download(
+            repo_id=self.repo_id,
+            filename="model.pt",
+            repo_type="model"
+        )
+        
+        self.config_path = hf_hub_download(
+            repo_id=self.repo_id,
+            filename="config.json",
+            repo_type="model"
+        )
+        
+        self.label_maps_path = hf_hub_download(
+            repo_id=self.repo_id,
+            filename="label_maps.json",
+            repo_type="model"
+        )
+        
+        self.metrics_path = self._try_download("metrics.json")
+        
+        self.config = self._load_json(self.config_path)
+        self.label_maps = self._load_json(self.label_maps_path)
+        self.metrics = self._load_json(self.metrics_path) if self.metrics_path else {}
+        
+        self.tasks = self.config.get("tasks") or list(self.label_maps.keys())
+        self.base_model_name = self.config.get("base_model_name") or self.config.get("model_name")
+        self.hidden_dim = self.config.get("hidden_dim", 512)
+        self.dropout = self.config.get("dropout", 0.2)
+        self.unfreeze_last_n_vision_layers = self.config.get("unfreeze_last_n_vision_layers", 0)
+        
+        if self.base_model_name is None:
+            self.base_model_name = "openai/clip-vit-base-patch32"
+        
+        self.task_num_classes = self._get_task_num_classes()
+        self.processor = CLIPImageProcessor.from_pretrained(self.base_model_name)
+        self.model = self._load_model()
+        self.model.eval()
+    
+    
+    def _try_download(self, filename):
+        try:
+            return hf_hub_download(
+                repo_id=self.repo_id,
+                filename=filename,
+                repo_type="model",
+            )
+        except Exception:
+            return None
+    
+    
+    def _load_json(self, path):
+        with open(path, "r", encoding="utf-8") as file:
+            return json.load(file)
+    
+    def _safe_torch_load(self, path):
+        try:
+            return torch.load(
+                path,
+                map_location=self.device,
+                weights_only=False
+            )
+        
+        except TypeError:
+            return torch.load(
+                path,
+                map_location=self.device
+            )
+    
+    def _get_task_num_classes(self):
+        if "task_num_classes" in self.config:
+            return {
+                task: int(value)
+                for task, value in self.config["task_num_classes"].items()
+            }
+        
+        task_num_classes = {}
+        for task in self.tasks:
+            task_num_classes[task] = len(self.label_maps[task]["label2id"])
+            
+        return task_num_classes
+    
+    def _load_model(self):
+        checkpoint = self._safe_torch_load(self.model_path)
+        
+        model = CLIPMultiTaskClassifier(
+            model_name=self.base_model_name,
+            task_num_classes=self.task_num_classes,
+            hidden_dim=self.hidden_dim,
+            droput=self.dropout,
+            unfreeze_last_n_vision_layers=self.unfreeze_last_n_vision_layers
+        )
+        
+        state_dict = checkpoint.get("model_state_dict", checkpoint)
+        model.load_state_dict(state_dict, strict=True)
+        model.to(self.device)
+        
+        return model
+        
+    
+    def _prepare_image(self, image):
+        if isinstance(image, Image.Image):
+            return image.convert("RGB")
+        
+        if isinstance(image, (str, Path)):
+            return Image.open(image).convert("RGB")
+        
+    
+    def predict(self, image, top_k=None):
+        top_k = top_k or self.top_k
+        image = self._prepare_image(image)
+        
+        inputs = self.processor(
+            images=image,
+            return_tensors="pt"
+        )
+        pixel_values = inputs["pixel_values"].to(self.device)
+        
+        if self.device == "cuda":
+            torch.cuda.synchronize()
+        
+        start_time = time.time()
+        
+        with torch.no_grad():
+            outputs = self.model(pixel_values)
+        
+        if self.device == "cuda":
+            torch.cuda.synchronize()
+        
+        end_time = time.time()
+        prediction = {}
+        simple_predictions = {}
+        
+        for task in self.tasks:
+            logits = outputs[task]
+            probs = torch.softmax(logits, dim=-1).squeeze(0)
+            
+            k = min(top_k, probs.shape[0])
+            top_probs, top_indices = torch.topk(
+                probs,
+                k=k
+            )
+            
+            top_predictions = []
+            for prob, idx in zip(top_probs, top_indices):
+                label = self.label_maps[task]["id2label"][str(int(idx.item()))] 
+                
+                top_predictions.append({
+                    "label": label,
+                    "confidence": float(prob.item()),
+                })
+            
+            prediction[task] = {
+                "label": top_predictions[0]["label"],
+                "confidence": top_predictions[0]["confidence"],
+                "top_3": top_predictions,
+            }
+
+            simple_predictions[task] = top_predictions[0]["label"]
+
+        catalog_output = generate_catalog_output(simple_predictions)
+
+        return {
+            "prediction": prediction,
+            "catalog_output": catalog_output,
+            "runtime": {
+                "device": self.device,
+                "inference_time_ms": float((end_time - start_time) * 1000),
+                "model": self.base_model_name,
+                "repo_id": self.repo_id,
+            },
+        }
+
+    def get_model_metrics(self):
+        return self.metrics

autocatalog/models/heads.py

@@ -0,0 +1,16 @@
+import torch.nn as nn
+
+class ClassificationHead(nn.Module):
+    def __init__(self, embedding_dim, num_classes, hidden_dim=512, dropout=0.2):
+        super().__init__()
+        
+        self.net = nn.Sequential(
+            nn.LayerNorm(embedding_dim),
+            nn.Linear(embedding_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_classes)
+        )
+    
+    def forward(self, x):
+        return self.net(x)

autocatalog/models/multitask_clip.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import CLIPModel
+from autocatalog.models.heads import ClassificationHead
+
+class CLIPMultiTaskClassifier(nn.Module):
+    def __init__(self, model_name, task_num_classes, hidden_dim=512, droput=0.2, unfreeze_last_n_vision_layers=0):
+        super().__init__()
+        self.clip = CLIPModel.from_pretrained(model_name)
+        
+        for param in self.clip.parameters():
+            param.requires_grad = False
+        
+        if unfreeze_last_n_vision_layers > 0:
+            vision_layer = self.clip.vision_model.encoder.layers
+            
+            for layer in vision_layer[-unfreeze_last_n_vision_layers:]:
+                for param in layer.parameters():
+                    param.requires_grad = True
+                
+            for param in self.clip.visual_projection.parameters():
+                param.requires_grad = True
+            
+            for param in self.clip.vision_model.post_layernorm.parameters():
+                param.requires_grad = True
+        
+        embedding_dim = self.clip.config.projection_dim
+        self.heads = nn.ModuleDict({
+            task : ClassificationHead(
+                embedding_dim=embedding_dim,
+                num_classes=num_classes,
+                hidden_dim=hidden_dim,
+                dropout=droput,
+            )
+            for task, num_classes in task_num_classes.items()
+        })
+
+    
+    def forward(self, pixel_values):
+        image_features = self.clip.get_image_features(pixel_values=pixel_values)
+        image_features = F.normalize(image_features, dim=-1)
+
+        return {
+            task : head(image_features)
+            for task, head in self.heads.items()
+        }

autocatalog/utils/config.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+import yaml
+
+def load_config(path):
+    config_path = Path(path)
+    if not config_path.exists():
+        return {}
+
+    with open(config_path, "r", encoding="utf-8") as file:
+        data = yaml.safe_load(file)
+
+    return data or {}

configs/config.yaml

@@ -0,0 +1,10 @@
+project:
+  name: AutoCatalogAI
+  version: v1
+
+model:
+  repo_id: mohsin416/autocatalogai-clip-multitask
+
+inference:
+  device: null
+  top_k: 3

notebooks/01_dataset_experiment.ipynb

@@ -0,0 +1,786 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "27ea1a10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q torch torchvision transformers datasets pillow pandas scikit-learn tqdm huggingface_hub matplotlib"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "d246862c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import random\n",
+    "import time\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "from PIL import Image\n",
+    "from tqdm.auto import tqdm\n",
+    "from datasets import load_dataset\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from transformers import CLIPModel, CLIPImageProcessor, get_cosine_schedule_with_warmup\n",
+    "\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, classification_report"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7e09b0d2",
+   "metadata": {},
+   "source": [
+    "_Config_"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "697fa289",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DATASET_NAME = \"ashraq/fashion-product-images-small\"\n",
+    "MODEL_NAME = \"openai/clip-vit-base-patch32\"\n",
+    "TASKS = [\n",
+    "    \"gender\",\n",
+    "    \"masterCategory\",\n",
+    "    \"subCategory\",\n",
+    "    \"articleType\",\n",
+    "    \"baseColour\",\n",
+    "    \"season\",\n",
+    "    \"usage\"\n",
+    "]\n",
+    "\n",
+    "SEED = 42\n",
+    "TRAIN_RATIO = 0.70\n",
+    "VAL_RATIO = 0.15\n",
+    "TEST_RATIO = 0.15\n",
+    "\n",
+    "BATCH_SIZE = 32\n",
+    "EPOCHS = 5\n",
+    "\n",
+    "HEAD_LR = 3e-4\n",
+    "BACKBONE_LR = 1e-5\n",
+    "WEIGHT_DECAY = 1e-2\n",
+    "\n",
+    "HIDDEN_DIM = 512\n",
+    "DROPOUT = 0.20\n",
+    "\n",
+    "UNFREEZE_LAST_N_VISION_LAYERS = 2\n",
+    "\n",
+    "USE_CLASS_WEIGHTS = True\n",
+    "USE_AMP = True\n",
+    "\n",
+    "MAX_GRAD_NORM = 1.0\n",
+    "EARLY_STOPPING_PATIENCE = 2\n",
+    "NUM_WORKERS = 2\n",
+    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "ef7f600c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Device: cuda\n",
+      "Model: openai/clip-vit-base-patch32\n"
+     ]
+    }
+   ],
+   "source": [
+    "OUTPUT_DIR = Path(\"artifacts/models/autocatalogai_clip\")\n",
+    "EVAL_DIR = Path(\"artifacts/evaluation\")\n",
+    "PLOT_DIR = Path(\"artifacts/plots\")\n",
+    "PROCESSED_DIR = Path(\"data/processed\")\n",
+    "\n",
+    "for directory in [OUTPUT_DIR, EVAL_DIR, PLOT_DIR, PROCESSED_DIR]:\n",
+    "    directory.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
+    "\n",
+    "print(\"Device:\", DEVICE)\n",
+    "print(\"Model:\", MODEL_NAME)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "d5242fdc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def set_seed(seed):\n",
+    "    random.seed(seed)\n",
+    "    np.random.seed(seed)\n",
+    "    torch.manual_seed(seed)\n",
+    "    \n",
+    "    if torch.cuda.is_available():\n",
+    "        torch.cuda.manual_seed_all(seed)\n",
+    "    torch.backends.cudnn.benchmark = True\n",
+    "\n",
+    "\n",
+    "set_seed(SEED)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f07779e1",
+   "metadata": {},
+   "source": [
+    "#### Load Full Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "280fb69a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:138: UserWarning: \n",
+      "Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'.\n",
+      "You are not authenticated with the Hugging Face Hub in this notebook.\n",
+      "If the error persists, please let us know by opening an issue on GitHub (https://github.com/huggingface/huggingface_hub/issues/new).\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f72208f4dc6348fc87d90c6c22769b0f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "README.md:   0%|          | 0.00/867 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n",
+      "WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ee563a2427bc4eb9b91f160c593b5477",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "data/train-00000-of-00002-6cff4c59f91661(…):   0%|          | 0.00/136M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2960cb54d2ba4573a7ef82894e897ad6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "data/train-00001-of-00002-bb459e5ac5f01e(…):   0%|          | 0.00/135M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7aa26aab19ff40e1a74b28de564323e2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Generating train split:   0%|          | 0/44072 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Dataset({\n",
+      "    features: ['id', 'gender', 'masterCategory', 'subCategory', 'articleType', 'baseColour', 'season', 'year', 'usage', 'productDisplayName', 'image'],\n",
+      "    num_rows: 44072\n",
+      "})\n",
+      "['id', 'gender', 'masterCategory', 'subCategory', 'articleType', 'baseColour', 'season', 'year', 'usage', 'productDisplayName', 'image']\n",
+      "Total rows: 44072\n"
+     ]
+    }
+   ],
+   "source": [
+    "raw_dataset = load_dataset(DATASET_NAME, split=\"train\")\n",
+    "\n",
+    "print(raw_dataset)\n",
+    "print(raw_dataset.column_names)\n",
+    "print(\"Total rows:\", len(raw_dataset))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5edb84b5",
+   "metadata": {},
+   "source": [
+    "#### Clean Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "bdf04d04",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "missing_columns = [task for task in TASKS if task not in raw_dataset.column_names]\n",
+    "\n",
+    "if \"image\" not in raw_dataset.column_names:\n",
+    "    raise ValueError(f\"Dataset must contain image column. Found: {raw_dataset.column_names}\")\n",
+    "\n",
+    "if missing_columns:\n",
+    "    raise ValueError(f\"Missing task columns: {missing_columns}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "721ec563",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "49bb2b38bea040a5896ceaf814335e21",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Filter:   0%|          | 0/44072 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Before cleaning: 44072\n",
+      "After cleaning: 44072\n"
+     ]
+    }
+   ],
+   "source": [
+    "def is_valid_row(row):\n",
+    "    for task in TASKS:\n",
+    "        value = row.get(task)\n",
+    "        if value is None:\n",
+    "            return False\n",
+    "        if str(value).strip() == \"\":\n",
+    "            return False\n",
+    "        \n",
+    "    return row.get(\"image\") is not None\n",
+    "\n",
+    "\n",
+    "clean_dataset = raw_dataset.filter(is_valid_row)\n",
+    "print(\"Before cleaning:\", len(raw_dataset))\n",
+    "print(\"After cleaning:\", len(clean_dataset))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "59e5f003",
+   "metadata": {},
+   "source": [
+    "#### Create Metadata DataFrame"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "17fecce2",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "  <div id=\"df-819ad3e2-5ed3-4458-823f-695ea26d18b3\" class=\"colab-df-container\">\n",
+       "    <div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>gender</th>\n",
+       "      <th>masterCategory</th>\n",
+       "      <th>subCategory</th>\n",
+       "      <th>articleType</th>\n",
+       "      <th>baseColour</th>\n",
+       "      <th>season</th>\n",
+       "      <th>usage</th>\n",
+       "      <th>id</th>\n",
+       "      <th>productDisplayName</th>\n",
+       "      <th>dataset_idx</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Men</td>\n",
+       "      <td>Apparel</td>\n",
+       "      <td>Topwear</td>\n",
+       "      <td>Shirts</td>\n",
+       "      <td>Navy Blue</td>\n",
+       "      <td>Fall</td>\n",
+       "      <td>Casual</td>\n",
+       "      <td>15970</td>\n",
+       "      <td>Turtle Check Men Navy Blue Shirt</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Men</td>\n",
+       "      <td>Apparel</td>\n",
+       "      <td>Bottomwear</td>\n",
+       "      <td>Jeans</td>\n",
+       "      <td>Blue</td>\n",
+       "      <td>Summer</td>\n",
+       "      <td>Casual</td>\n",
+       "      <td>39386</td>\n",
+       "      <td>Peter England Men Party Blue Jeans</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>Women</td>\n",
+       "      <td>Accessories</td>\n",
+       "      <td>Watches</td>\n",
+       "      <td>Watches</td>\n",
+       "      <td>Silver</td>\n",
+       "      <td>Winter</td>\n",
+       "      <td>Casual</td>\n",
+       "      <td>59263</td>\n",
+       "      <td>Titan Women Silver Watch</td>\n",
+       "      <td>2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Men</td>\n",
+       "      <td>Apparel</td>\n",
+       "      <td>Bottomwear</td>\n",
+       "      <td>Track Pants</td>\n",
+       "      <td>Black</td>\n",
+       "      <td>Fall</td>\n",
+       "      <td>Casual</td>\n",
+       "      <td>21379</td>\n",
+       "      <td>Manchester United Men Solid Black Track Pants</td>\n",
+       "      <td>3</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>Men</td>\n",
+       "      <td>Apparel</td>\n",
+       "      <td>Topwear</td>\n",
+       "      <td>Tshirts</td>\n",
+       "      <td>Grey</td>\n",
+       "      <td>Summer</td>\n",
+       "      <td>Casual</td>\n",
+       "      <td>53759</td>\n",
+       "      <td>Puma Men Grey T-shirt</td>\n",
+       "      <td>4</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>\n",
+       "    <div class=\"colab-df-buttons\">\n",
+       "      \n",
+       "  <div class=\"colab-df-container\">\n",
+       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-819ad3e2-5ed3-4458-823f-695ea26d18b3')\"\n",
+       "            title=\"Convert this dataframe to an interactive table.\"\n",
+       "            style=\"display:none;\">\n",
+       "      \n",
+       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
+       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
+       "  </svg>\n",
+       "    </button>\n",
+       "    \n",
+       "  <style>\n",
+       "    .colab-df-container {\n",
+       "      display:flex;\n",
+       "      gap: 12px;\n",
+       "    }\n",
+       "\n",
+       "    .colab-df-convert {\n",
+       "      background-color: #E8F0FE;\n",
+       "      border: none;\n",
+       "      border-radius: 50%;\n",
+       "      cursor: pointer;\n",
+       "      display: none;\n",
+       "      fill: #1967D2;\n",
+       "      height: 32px;\n",
+       "      padding: 0 0 0 0;\n",
+       "      width: 32px;\n",
+       "    }\n",
+       "\n",
+       "    .colab-df-convert:hover {\n",
+       "      background-color: #E2EBFA;\n",
+       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
+       "      fill: #174EA6;\n",
+       "    }\n",
+       "\n",
+       "    .colab-df-buttons div {\n",
+       "      margin-bottom: 4px;\n",
+       "    }\n",
+       "\n",
+       "    [theme=dark] .colab-df-convert {\n",
+       "      background-color: #3B4455;\n",
+       "      fill: #D2E3FC;\n",
+       "    }\n",
+       "\n",
+       "    [theme=dark] .colab-df-convert:hover {\n",
+       "      background-color: #434B5C;\n",
+       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
+       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
+       "      fill: #FFFFFF;\n",
+       "    }\n",
+       "  </style>\n",
+       "\n",
+       "    <script>\n",
+       "      const buttonEl =\n",
+       "        document.querySelector('#df-819ad3e2-5ed3-4458-823f-695ea26d18b3 button.colab-df-convert');\n",
+       "      buttonEl.style.display =\n",
+       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
+       "\n",
+       "      async function convertToInteractive(key) {\n",
+       "        const element = document.querySelector('#df-819ad3e2-5ed3-4458-823f-695ea26d18b3');\n",
+       "        const dataTable =\n",
+       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
+       "                                                    [key], {});\n",
+       "        if (!dataTable) return;\n",
+       "\n",
+       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
+       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
+       "          + ' to learn more about interactive tables.';\n",
+       "        element.innerHTML = '';\n",
+       "        dataTable['output_type'] = 'display_data';\n",
+       "        await google.colab.output.renderOutput(dataTable, element);\n",
+       "        const docLink = document.createElement('div');\n",
+       "        docLink.innerHTML = docLinkHtml;\n",
+       "        element.appendChild(docLink);\n",
+       "      }\n",
+       "    </script>\n",
+       "  </div>\n",
+       "  \n",
+       "    </div>\n",
+       "  </div>\n",
+       "  "
+      ],
+      "text/plain": [
+       "  gender masterCategory subCategory  articleType baseColour  season   usage  \\\n",
+       "0    Men        Apparel     Topwear       Shirts  Navy Blue    Fall  Casual   \n",
+       "1    Men        Apparel  Bottomwear        Jeans       Blue  Summer  Casual   \n",
+       "2  Women    Accessories     Watches      Watches     Silver  Winter  Casual   \n",
+       "3    Men        Apparel  Bottomwear  Track Pants      Black    Fall  Casual   \n",
+       "4    Men        Apparel     Topwear      Tshirts       Grey  Summer  Casual   \n",
+       "\n",
+       "      id                             productDisplayName  dataset_idx  \n",
+       "0  15970               Turtle Check Men Navy Blue Shirt            0  \n",
+       "1  39386             Peter England Men Party Blue Jeans            1  \n",
+       "2  59263                       Titan Women Silver Watch            2  \n",
+       "3  21379  Manchester United Men Solid Black Track Pants            3  \n",
+       "4  53759                          Puma Men Grey T-shirt            4  "
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "metadata = {}\n",
+    "for_col = []\n",
+    "extra_columns = [\"id\", \"productDisplayName\"]\n",
+    "\n",
+    "for task in TASKS:\n",
+    "    metadata[task] = [str(value).strip() for value in clean_dataset[task]]\n",
+    "\n",
+    "\n",
+    "for col in extra_columns:\n",
+    "    if col in clean_dataset.column_names:\n",
+    "        metadata[col] = clean_dataset[col]\n",
+    "        for_col.append(col)\n",
+    "\n",
+    "df = pd.DataFrame(metadata)\n",
+    "df[\"dataset_idx\"] = np.arange(len(clean_dataset))\n",
+    "\n",
+    "df.head()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d55ac03",
+   "metadata": {},
+   "source": [
+    "#### Label Distribution"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "5fe58be4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "label_distribution = {}\n",
+    "\n",
+    "for task in TASKS:\n",
+    "    counts = df[task].value_counts().to_dict()\n",
+    "    label_distribution[task] = counts\n",
+    "    \n",
+    "with open(EVAL_DIR / \"label_distribution.json\", \"w\", encoding=\"utf-8\") as f:\n",
+    "    json.dump(label_distribution, f, indent=2, ensure_ascii=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "f77e1f0d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'dataset_name': 'ashraq/fashion-product-images-small',\n",
+       " 'total_clean_samples': 44072,\n",
+       " 'tasks': ['gender',\n",
+       "  'masterCategory',\n",
+       "  'subCategory',\n",
+       "  'articleType',\n",
+       "  'baseColour',\n",
+       "  'season',\n",
+       "  'usage'],\n",
+       " 'num_classes': {'gender': 5,\n",
+       "  'masterCategory': 7,\n",
+       "  'subCategory': 45,\n",
+       "  'articleType': 141,\n",
+       "  'baseColour': 46,\n",
+       "  'season': 4,\n",
+       "  'usage': 8}}"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "summary = {\n",
+    "    \"dataset_name\": DATASET_NAME,\n",
+    "    \"total_clean_samples\": len(df),\n",
+    "    \"tasks\": TASKS,\n",
+    "    \"num_classes\": {\n",
+    "        task: int(df[task].nunique())\n",
+    "        for task in TASKS\n",
+    "    }\n",
+    "}\n",
+    "\n",
+    "with open(EVAL_DIR / \"dataset_summary.json\", \"w\", encoding=\"utf-8\") as f:\n",
+    "    json.dump(summary, f, indent=2, ensure_ascii=False)\n",
+    "\n",
+    "summary"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b73fb107",
+   "metadata": {},
+   "source": [
+    "#### Train / Validation / Test Split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "ff328f23",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_safe_stratify_labels(series):\n",
+    "    counts = series.value_counts()\n",
+    "    return series.apply(lambda x: x if counts[x] >= 2 else \"__rare__\")\n",
+    "\n",
+    "stratify_labels = make_safe_stratify_labels(df[\"articleType\"])\n",
+    "all_indices = df.index.to_numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "c876c207",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_idx, temp_idx = train_test_split(\n",
+    "    all_indices,\n",
+    "    test_size=0.30,\n",
+    "    random_state=SEED,\n",
+    "    stratify=stratify_labels\n",
+    ")\n",
+    "\n",
+    "temp_df = df.loc[temp_idx].copy()\n",
+    "temp_stratify_labels = make_safe_stratify_labels(temp_df[\"articleType\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "0ffaad6a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_idx, test_idx = train_test_split(\n",
+    "    temp_idx,\n",
+    "    test_size=0.50,\n",
+    "    random_state=SEED,\n",
+    "    stratify=temp_stratify_labels\n",
+    ")\n",
+    "\n",
+    "train_df = df.loc[train_idx].copy()\n",
+    "val_df = df.loc[val_idx].copy()\n",
+    "test_df = df.loc[test_idx].copy()\n",
+    "\n",
+    "train_df[\"split\"] = \"train\"\n",
+    "val_df[\"split\"] = \"validation\"\n",
+    "test_df[\"split\"] = \"test\"\n",
+    "\n",
+    "train_df.to_csv(PROCESSED_DIR / \"train.csv\", index=False)\n",
+    "val_df.to_csv(PROCESSED_DIR / \"val.csv\", index=False)\n",
+    "test_df.to_csv(PROCESSED_DIR / \"test.csv\", index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "0bd769f2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train: 30850 0.7\n",
+      "Validation: 6611 0.15\n",
+      "Test: 6611 0.15\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"Train:\", len(train_df), round(len(train_df) / len(df), 3))\n",
+    "print(\"Validation:\", len(val_df), round(len(val_df) / len(df), 3))\n",
+    "print(\"Test:\", len(test_df), round(len(test_df) / len(df), 3))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1b2990da",
+   "metadata": {},
+   "source": [
+    "#### Build HF Split Datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "c42e21e5",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(30850, 6611, 6611)"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "train_hf_dataset = clean_dataset.select(train_df[\"dataset_idx\"].tolist())\n",
+    "val_hf_dataset = clean_dataset.select(val_df[\"dataset_idx\"].tolist())\n",
+    "test_hf_dataset = clean_dataset.select(test_df[\"dataset_idx\"].tolist())\n",
+    "\n",
+    "len(train_hf_dataset), len(val_hf_dataset), len(test_hf_dataset)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

requirements.txt

@@ -1,3 +1,7 @@
-altair
-pandas
-streamlit
+streamlit
+torch
+torchvision
+transformers
+huggingface_hub
+Pillow
+PyYAML

setup.py

@@ -0,0 +1,9 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="autocatalogai",
+    version="0.1.0",
+    author="Md Mohsin",
+    author_email="siam.mohsin2005@gmail.com",
+    packages=find_packages()
+)

src/streamlit_app.py

@@ -1,40 +0,0 @@
-import altair as alt
-import numpy as np
-import pandas as pd
-import streamlit as st
-
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))

template.py

@@ -0,0 +1,80 @@
+import os
+from pathlib import Path
+
+project_name = "AutoCatalogAI"
+list_of_files = [
+    f"{project_name}/app/app.py",
+    f"{project_name}/app/inference.py",
+    f"{project_name}/app/templates/index.html",
+
+    f"{project_name}/autocatalog/__init__.py",
+
+    f"{project_name}/autocatalog/data/__init__.py",
+    f"{project_name}/autocatalog/data/dataset.py",
+    f"{project_name}/autocatalog/data/preprocessing.py",
+
+    f"{project_name}/autocatalog/models/__init__.py",
+    f"{project_name}/autocatalog/models/multitask_clip.py",
+    f"{project_name}/autocatalog/models/heads.py",
+
+    f"{project_name}/autocatalog/training/__init__.py",
+    f"{project_name}/autocatalog/training/train.py",
+    f"{project_name}/autocatalog/training/losses.py",
+
+    f"{project_name}/autocatalog/evaluation/__init__.py",
+    f"{project_name}/autocatalog/evaluation/evaluate.py",
+    f"{project_name}/autocatalog/evaluation/metrics.py",
+    f"{project_name}/autocatalog/evaluation/error_analysis.py",
+
+    f"{project_name}/autocatalog/inference/__init__.py",
+    f"{project_name}/autocatalog/inference/predictor.py",
+    f"{project_name}/autocatalog/inference/catalog_generator.py",
+
+    f"{project_name}/autocatalog/utils/__init__.py",
+    f"{project_name}/autocatalog/utils/config.py",
+    f"{project_name}/autocatalog/utils/logger.py",
+    f"{project_name}/autocatalog/utils/seed.py",
+
+    f"{project_name}/configs/config.yaml",
+
+    f"{project_name}/data/processed/train.csv",
+    f"{project_name}/data/processed/val.csv",
+    f"{project_name}/data/processed/test.csv",
+
+    f"{project_name}/artifacts/models/.gitkeep",
+    f"{project_name}/artifacts/evaluation/.gitkeep",
+    f"{project_name}/artifacts/plots/.gitkeep",
+    f"{project_name}/artifacts/examples/.gitkeep",
+
+    f"{project_name}/notebooks/01_dataset_experiment.ipynb",
+
+    f"{project_name}/scripts/prepare_dataset.py",
+    f"{project_name}/scripts/train_baseline.py",
+    f"{project_name}/scripts/train_multitask_clip.py",
+    f"{project_name}/scripts/evaluate_model.py",
+    f"{project_name}/scripts/predict_image.py",
+
+    f"{project_name}/tests/test_dataset.py",
+    f"{project_name}/tests/test_model.py",
+    f"{project_name}/tests/test_inference.py",
+
+    f"{project_name}/.gitignore",
+    f"{project_name}/README.md",
+    f"{project_name}/model_card.md",
+    f"{project_name}/requirements.txt",
+    f"{project_name}/Dockerfile",
+    f"{project_name}/setup.py",
+]
+
+for filepath in list_of_files:
+    filepath = Path(filepath)
+    filedir, filename = os.path.split(filepath)
+
+    if filedir:
+        os.makedirs(filedir, exist_ok=True)
+
+    if not filepath.exists():
+        filepath.touch()
+        print(f"Created: {filepath}")
+    else:
+        print(f"Already exists: {filepath}")