rbx-proofreader 1.0.1__py3-none-any.whl → 1.1.1__py3-none-any.whl

proofreader/core/config.py

@@ -5,13 +5,17 @@ from pathlib import Path
 # --- BASE PATHS ---
 # Resolves to the 'proofreader' root directory
 BASE_DIR = Path(__file__).resolve().parent.parent.parent
+BASE_URL = "https://github.com/lucacrose/proofreader"
 
 # --- ASSETS & MODELS ---
 ASSETS_PATH = BASE_DIR / "assets"
 MODEL_PATH = ASSETS_PATH / "weights" / "yolo.pt"
-DB_PATH = ASSETS_PATH / "db.json"
-CACHE_PATH = ASSETS_PATH / "embedding_bank.pt"
+DB_PATH = ASSETS_PATH / "item_database.json"
+CACHE_PATH = ASSETS_PATH / "item_embeddings_bank.pt"
 THUMBNAILS_DIR = ASSETS_PATH / "thumbnails"
+TRAIN_THUMBNAILS_DIR = ASSETS_PATH / "train_data"
+CLASS_MAP_PATH = ASSETS_PATH / "class_mapping.json"
+CLIP_BEST_PATH = ASSETS_PATH / "weights" / "clip.pt"
 
 # --- TRAINING & EMULATOR ---
 TRAIN_DIR = BASE_DIR / "proofreader" / "train"
@@ -26,11 +30,11 @@ DEFAULT_TEMPLATE = TEMPLATES_DIR / "trade_ui.html"
 
 # --- HYPERPARAMETERS (Training Settings) ---
 TRAINING_CONFIG = {
-    "epochs": 100,             # Number of times the model sees the whole dataset
+    "epochs": 240,             # Number of times the model sees the whole dataset
     "batch_size": 16,          # Number of images processed at once
     "img_size": 640,           # Standard YOLO resolution
-    "patience": 10,            # Stop early if no improvement for 10 epochs
-    "close_mosaic_epochs": 10  # Disable mosaic augmentation for the last N epochs
+    "patience": 20,            # Stop early if no improvement for 20 epochs
+    "close_mosaic_epochs": 32  # Disable mosaic augmentation for the last N epochs
 }
 
 # --- AUGMENTER PROBABILITIES AND GENERATOR SETTINGS ---
@@ -82,7 +86,7 @@ AUGMENTER_CONFIG = {
 
 # Robustness Thresholds
 FUZZY_MATCH_CONFIDENCE_THRESHOLD = 60.0
-VISUAL_MATCH_THRESHOLD = 0.88
+CERTAIN_VISUAL_CONF = 0.995
 
 # --- HARDWARE SETTINGS ---
 # Automatically detects if a GPU is available for faster training

proofreader/core/matcher.py

@@ -1,34 +1,56 @@
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import numpy as np
+import json
 import cv2
 from PIL import Image
-from typing import Dict, List, Any
-from .schema import TradeLayout
-from proofreader.core.config import VISUAL_MATCH_THRESHOLD
+from torchvision import transforms
+from transformers import CLIPVisionModelWithProjection
+from typing import List
+from .schema import TradeLayout, ResolvedItem
+
+class CLIPItemEmbedder(nn.Module):
+    def __init__(self, num_classes, model_id="openai/clip-vit-base-patch32"):
+        super().__init__()
+        self.vision_encoder = CLIPVisionModelWithProjection.from_pretrained(model_id)
+        self.item_prototypes = nn.Embedding(num_classes, 512)
+        self.logit_scale = nn.Parameter(torch.ones([]) * 2.659)
+
+    def forward(self, pixel_values):
+        outputs = self.vision_encoder(pixel_values=pixel_values)
+        return F.normalize(outputs.image_embeds, p=2, dim=-1)
 
 class VisualMatcher:
-    def __init__(self, embedding_bank: Dict[str, np.ndarray], item_db: List[dict], clip_processor: Any, clip_model: Any, device: str = "cuda"):
+    def __init__(self, weights_path: str, mapping_path: str, item_db: List[dict], device: str = "cuda"):
         self.device = device
-        self.bank = embedding_bank
-        self.item_db = item_db
-        self.clip_processor = clip_processor
-        self.clip_model = clip_model
 
-        self.name_to_id = {str(i["name"]).lower().strip(): i["id"] for i in item_db}
+        with open(mapping_path, "r") as f:
+            self.class_to_idx = json.load(f)
+        self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}
+
         self.id_to_name = {str(i["id"]): i["name"] for i in item_db}
+        self.name_to_id = {str(i["name"]).lower().strip(): i["id"] for i in item_db}
 
-        self.bank_names = list(embedding_bank.keys())
-        self.bank_tensor = torch.stack([embedding_bank[name] for name in self.bank_names]).to(self.device)
-        self.bank_tensor = torch.nn.functional.normalize(self.bank_tensor, dim=1)
+        num_classes = len(self.class_to_idx)
+        self.model = CLIPItemEmbedder(num_classes).to(self.device)
+        self.model.load_state_dict(torch.load(weights_path, map_location=self.device))
+        self.model.eval()
 
-    def _get_id_from_name(self, name: str) -> str:
-        item = next((i for i in self.item_db if i["name"] == name), None)
-        return item["id"] if item else 0
+        with torch.inference_mode():
+            self.bank_tensor = F.normalize(self.model.item_prototypes.weight, p=2, dim=-1)
+        
+        self.preprocess = transforms.Compose([
+            transforms.Resize((224, 224), interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), 
+                                 (0.26862954, 0.26130258, 0.27577711)),
+        ])
 
-    def match_item_visuals(self, image: np.ndarray, layout: TradeLayout, similarity_threshold: float = VISUAL_MATCH_THRESHOLD):
-        items_to_process = []
+    def match_item_visuals(self, image: np.ndarray, layout: TradeLayout):
+        items_to_process: List[ResolvedItem] = []
         crops = []
-        
+
         for side in (layout.outgoing.items, layout.incoming.items):
             for item in side:
                 if item.thumb_box:
@@ -36,32 +58,32 @@ class VisualMatcher:
                     crop = image[y1:y2, x1:x2]
                     if crop.size > 0:
                         pil_img = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
-                        crops.append(pil_img)
+                        processed_crop = self.preprocess(pil_img)
+                        crops.append(processed_crop)
                         items_to_process.append(item)
 
         if not crops:
             return
         
-        inputs = self.clip_processor(images=crops, return_tensors="pt", padding=True).to(self.device)
+        batch_tensor = torch.stack(crops).to(self.device)
         
-        with torch.no_grad():
-            query_features = self.clip_model.get_image_features(**inputs)
-            query_features = torch.nn.functional.normalize(query_features, dim=1)
-            similarities = torch.matmul(query_features, self.bank_tensor.T)
-            best_scores, best_indices = torch.max(similarities, dim=1)
+        with torch.inference_mode():
+            query_features = self.model(batch_tensor)
+
+            logits = query_features @ self.bank_tensor.t() * self.model.logit_scale.exp()
+            topk_scores, topk_indices = logits.topk(k=5, dim=1)
+
+            probs = F.softmax(topk_scores.float(), dim=1)
+            
+            best_idx_in_topk = probs.argmax(dim=1)
+            best_indices = topk_indices[torch.arange(len(topk_indices)), best_idx_in_topk]
+            best_probs = probs[torch.arange(len(probs)), best_idx_in_topk]
+
         
         for i, item in enumerate(items_to_process):
-            visual_match_val = self.bank_names[best_indices[i]]
-            visual_conf = best_scores[i].item()
+            visual_idx = best_indices[i].item()
 
-            is_ocr_valid = item.name.lower().strip() in self.name_to_id if item.name else False
+            visual_match_id_str = self.idx_to_class[visual_idx]
 
-            if (not is_ocr_valid or visual_conf > 0.95) and visual_conf >= similarity_threshold:
-                if str(visual_match_val).isdigit():
-                    item.id = int(visual_match_val)
-                    item.name = self.id_to_name.get(str(visual_match_val), "Unknown Item")
-                else:
-                    item.name = visual_match_val
-                    item.id = self._get_id_from_name(visual_match_val)
-            else:
-                item.id = self._get_id_from_name(item.name)
+            item.visual_id = int(visual_match_id_str)
+            item.visual_conf = float(best_probs[i].item())

proofreader/core/ocr.py

@@ -3,17 +3,13 @@ import easyocr
 import numpy as np
 import re
 from rapidfuzz import process, utils
-from .schema import Box, TradeLayout, TradeSide
+from .schema import TradeLayout
 from proofreader.core.config import FUZZY_MATCH_CONFIDENCE_THRESHOLD, OCR_LANGUAGES, OCR_USE_GPU
 
 class OCRReader:
     def __init__(self, item_list, languages=OCR_LANGUAGES, gpu=OCR_USE_GPU):
         self.reader = easyocr.Reader(languages, gpu=gpu)
-
-        self.item_names = []
-
-        for item in item_list:
-            self.item_names.append(item["name"])
+        self.item_names = [item["name"] for item in item_list]
 
     def _fuzzy_match_name(self, raw_text: str, threshold: float = FUZZY_MATCH_CONFIDENCE_THRESHOLD) -> str:
         if not raw_text or len(raw_text) < 2:
@@ -32,48 +28,65 @@ class OCRReader:
 
     def _clean_robux_text(self, raw_text: str) -> int:
         cleaned = raw_text.upper().strip()
-
         substitutions = {
             ',': '', '.': '', ' ': '',
             'S': '5', 'O': '0', 'I': '1', 
             'L': '1', 'B': '8', 'G': '6'
         }
-
         for char, sub in substitutions.items():
             cleaned = cleaned.replace(char, sub)
         
         digits = re.findall(r'\d+', cleaned)
-
         return int("".join(digits)) if digits else 0
 
-    def _get_text_from_box(self, image: np.ndarray, box: Box, is_robux: bool = False) -> str:
-        x1, y1, x2, y2 = box.coords
+    def process_layout(self, image: np.ndarray, layout: TradeLayout, skip_if=None):
+        all_items = layout.outgoing.items + layout.incoming.items
+        crops = []
+        target_refs = []
+        STD_H = 64 
 
-        crop = image[max(0, y1-2):y2+2, max(0, x1-2):x2+2]
-        
-        if crop.size == 0:
-            return ""
-        
-        gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
-        
-        if is_robux:
-            gray = cv2.resize(gray, None, fx=3, fy=3, interpolation=cv2.INTER_CUBIC)
-            results = self.reader.readtext(gray, allowlist="0123456789,S ")
-        else:
-            results = self.reader.readtext(gray)
-            
-        return " ".join([res[1] for res in results]).strip()
+        for item in all_items:
+            if skip_if and skip_if(item):
+                continue
 
-    def process_side(self, image: np.ndarray, side: TradeSide):
-        for item in side.items:
             if item.name_box:
-                raw_name = self._get_text_from_box(image, item.name_box)
-                item.name = self._fuzzy_match_name(raw_name)
+                x1, y1, x2, y2 = item.name_box.coords
+                crop = image[max(0, y1-2):y2+2, max(0, x1-2):x2+2]
+                if crop.size > 0:
+                    gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
+                    h, w = gray.shape
+                    new_w = int(w * (STD_H / h))
+                    resized = cv2.resize(gray, (new_w, STD_H), interpolation=cv2.INTER_LINEAR)
+                    crops.append(resized)
+                    target_refs.append({'type': 'item', 'obj': item})
+        
+        for side in [layout.outgoing, layout.incoming]:
+            if side.robux and side.robux.value_box:
+                x1, y1, x2, y2 = side.robux.value_box.coords
+                crop = image[max(0, y1-2):y2+2, max(0, x1-2):x2+2]
+                if crop.size > 0:
+                    gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
+                    h, w = gray.shape
+                    new_w = int(w * (STD_H / h))
+                    resized = cv2.resize(gray, (new_w, STD_H), interpolation=cv2.INTER_LINEAR)
+                    crops.append(resized)
+                    target_refs.append({'type': 'robux', 'obj': side.robux})
 
-        if side.robux and side.robux.value_box:
-            raw_val = self._get_text_from_box(image, side.robux.value_box, is_robux=True)
-            side.robux.value = self._clean_robux_text(raw_val)
+        if not crops:
+            return
+        
+        max_w = max(c.shape[1] for c in crops)
+        padded_crops = [cv2.copyMakeBorder(c, 0, 0, 0, max_w - c.shape[1], cv2.BORDER_CONSTANT, value=0) for c in crops]
+
+        batch_results = self.reader.readtext_batched(padded_crops, batch_size=len(padded_crops))
 
-    def process_layout(self, image: str, layout: TradeLayout):
-        self.process_side(image, layout.outgoing)
-        self.process_side(image, layout.incoming)
+        for i, res in enumerate(batch_results):
+            raw_text = " ".join([text_info[1] for text_info in res]).strip()
+            conf = np.mean([text_info[2] for text_info in res]) if res else 0.0
+            
+            target = target_refs[i]
+            if target['type'] == 'item':
+                target['obj'].text_name = raw_text
+                target['obj'].text_conf = float(conf)
+            else:
+                target['obj'].value = self._clean_robux_text(raw_text)

proofreader/core/schema.py

@@ -15,6 +15,14 @@ class ResolvedItem:
     thumb_box: Optional[Box] = None
     name_box: Optional[Box] = None
 
+    visual_id: int = -1
+    visual_conf: float = 0
+
+    text_name: str = ""
+    text_conf: float = 0
+
+    _finalized: bool = False
+
 @dataclass
 class ResolvedRobux:
     value: int = 0

proofreader/main.py

@@ -9,7 +9,8 @@ from .core.detector import TradeDetector
 from .core.resolver import SpatialResolver
 from .core.ocr import OCRReader
 from .core.matcher import VisualMatcher
-from .core.config import DB_PATH, CACHE_PATH, MODEL_PATH, DEVICE
+from .core.config import DB_PATH, MODEL_PATH, DEVICE, CLASS_MAP_PATH, CLIP_BEST_PATH, BASE_URL, CERTAIN_VISUAL_CONF
+from .core.schema import ResolvedItem
 
 class TradeEngine:
     def __init__(self):
@@ -28,38 +29,32 @@ class TradeEngine:
 
         self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(self.device)
         self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", use_fast=True)
-        
+
         with open(DB_PATH, "r") as f:
             item_db = json.load(f)
         
-        cache_data = torch.load(CACHE_PATH, weights_only=False)['embeddings']
-        self.embeddings = {k: torch.tensor(v).to(self.device) for k, v in cache_data.items()}
-
         self.detector = TradeDetector(MODEL_PATH)
         self.resolver = SpatialResolver()
-        
         self.reader = OCRReader(item_db)
-        
+
         self.matcher = VisualMatcher(
-            embedding_bank=self.embeddings,
             item_db=item_db,
-            clip_processor=self.clip_processor,
-            clip_model=self.clip_model,
+            weights_path=CLIP_BEST_PATH,
+            mapping_path=CLASS_MAP_PATH,
             device=self.device
         ) 
 
     def _ensure_assets(self):
-        BASE_URL = "https://github.com/lucacrose/proofreader"
-        
         assets = {
-            DB_PATH: f"{BASE_URL}/releases/download/v1.0.0/db.json",
-            CACHE_PATH: f"{BASE_URL}/releases/download/v1.0.0/embedding_bank.pt",
-            MODEL_PATH: f"{BASE_URL}/releases/download/v1.0.0/yolo.pt"
+            DB_PATH: f"{BASE_URL}/releases/download/v1.1.0/item_database.json",
+            MODEL_PATH: f"{BASE_URL}/releases/download/v1.1.0/yolo.pt",
+            CLIP_BEST_PATH: f"{BASE_URL}/releases/download/v1.1.0/clip.pt",
+            CLASS_MAP_PATH: f"{BASE_URL}/releases/download/v1.1.0/class_mapping.json"
         }
 
         for path, url in assets.items():
             if not path.exists():
-                print(f"📦 {path.name} missing. Downloading from published release...")
+                print(f"📦 {path.name} missing. Downloading from latest release...")
                 self._download_file(url, path)
 
     def _download_file(self, url, dest_path):
@@ -73,17 +68,73 @@ class TradeEngine:
                 f.write(chunk)
                 pbar.update(len(chunk))
 
+    def _final_judge(self, item: ResolvedItem):
+        if getattr(item, "_finalized", False):
+            return
+        
+        v_id = item.visual_id
+        v_conf = item.visual_conf
+
+        ocr_name_raw = item.text_name.lower().strip()
+        ocr_id_direct = self.matcher.name_to_id.get(ocr_name_raw)
+        ocr_conf = item.text_conf / 100.0 if item.text_conf > 1 else item.text_conf
+
+        if v_id != -1 and v_id == ocr_id_direct:
+            item.id = v_id
+            item.name = self.matcher.id_to_name.get(str(v_id))
+            return
+        
+        if v_conf > 0.85:
+            item.id = v_id
+            item.name = self.matcher.id_to_name.get(str(v_id))
+            return
+        
+        if ocr_conf > 0.85 and ocr_id_direct:
+            item.id = ocr_id_direct
+            item.name = self.matcher.id_to_name.get(str(ocr_id_direct))
+            return
+        
+        if len(ocr_name_raw) > 2:
+            fuzzy_name = self.reader._fuzzy_match_name(ocr_name_raw)
+            fuzzy_id = self.matcher.name_to_id.get(fuzzy_name.lower())
+            
+            if fuzzy_id:
+                item.id = int(fuzzy_id)
+                item.name = fuzzy_name
+                return
+        
+        if v_conf >= ocr_conf and v_id != -1:
+            item.id = v_id
+            item.name = self.matcher.id_to_name.get(str(v_id))
+        elif ocr_id_direct:
+            item.id = ocr_id_direct
+            item.name = self.matcher.id_to_name.get(str(ocr_id_direct))
+
     def process_image(self, image_path: str, conf_threshold: float) -> dict:
         if not os.path.exists(image_path):
             raise FileNotFoundError(f"Image not found: {image_path}")
         
         boxes = self.detector.detect(image_path, conf_threshold)
         layout = self.resolver.resolve(boxes)
-
         image = cv2.imread(image_path)
 
-        self.reader.process_layout(image, layout)
-
         self.matcher.match_item_visuals(image, layout)
 
+        for side in [layout.outgoing, layout.incoming]:
+            for item in side.items:
+                if item.visual_id != -1 and item.visual_conf >= CERTAIN_VISUAL_CONF:
+                    item.id = item.visual_id
+                    item.name = self.matcher.id_to_name.get(str(item.visual_id), "Unknown")
+                    item._finalized = True
+
+        self.reader.process_layout(
+            image,
+            layout,
+            skip_if=lambda item: getattr(item, "_finalized", False)
+        )
+
+        for side in [layout.outgoing, layout.incoming]:
+            for item in side.items:
+                self._final_judge(item)
+
         return layout.to_dict()

proofreader/train/clip_trainer.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from transformers import CLIPVisionModelWithProjection
+from tqdm import tqdm
+from torch.amp import GradScaler, autocast
+from proofreader.core.config import CLASS_MAP_PATH, CLIP_BEST_PATH, DATASET_ROOT
+import os
+import json
+import numpy as np
+import random
+
+MODEL_ID = "openai/clip-vit-base-patch32"
+EPOCHS = 10
+BATCH_SIZE = 48 
+LEARNING_RATE = 1e-5
+EMBEDDING_DIM = 512
+WEIGHT_DECAY = 0.1
+PATIENCE = 3        # Stop if no improvement for 3 epochs
+MIN_DELTA = 0.1     # Minimum % improvement to be considered "better"
+
+def set_seed(seed: int = 42):
+    random.seed(seed)
+    np.random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+class CLIPItemEmbedder(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.vision_encoder = CLIPVisionModelWithProjection.from_pretrained(MODEL_ID)
+        self.item_prototypes = nn.Embedding(num_classes, EMBEDDING_DIM)
+        self.logit_scale = nn.Parameter(torch.ones([]) * 2.659)
+
+    def forward(self, pixel_values, item_ids):
+        outputs = self.vision_encoder(pixel_values=pixel_values)
+        image_embeds = outputs.image_embeds 
+        
+        label_embeds = self.item_prototypes(item_ids)
+        label_embeds = F.normalize(label_embeds, p=2, dim=-1)
+        
+        return image_embeds, label_embeds, self.logit_scale.exp()
+
+class EarlyStopper:
+    def __init__(self, patience=3, min_delta=0.05):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.counter = 0
+        self.best_accuracy = 0
+        self.best_state = None
+
+    def check(self, current_accuracy, model):
+        if current_accuracy > (self.best_accuracy + self.min_delta):
+            self.best_accuracy = current_accuracy
+            self.best_state = getattr(model, "_orig_mod", model).state_dict()
+            self.counter = 0
+            return False, True
+        else:
+            self.counter += 1
+            return (self.counter >= self.patience), False
+
+def get_transforms():
+    return transforms.Compose([
+        transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
+        # Resolution Crush
+        transforms.RandomApply([
+            transforms.RandomChoice([transforms.Resize(128), transforms.Resize(64)]),
+            transforms.Resize(224),
+        ], p=0.3),
+        # Gaussian Blur
+        transforms.RandomApply([
+            transforms.GaussianBlur(kernel_size=(3, 5), sigma=(0.1, 2.0))
+        ], p=0.2),
+        transforms.CenterCrop(224),
+        transforms.RandomHorizontalFlip(),
+        transforms.ColorJitter(0.3, 0.3, 0.3),
+        transforms.ToTensor(),
+        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), 
+                             (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+def train_clip():
+    set_seed(1)
+
+    torch.backends.cudnn.benchmark = True 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    dataset_path = f"{DATASET_ROOT}/classification"
+    full_dataset = datasets.ImageFolder(root=dataset_path, transform=get_transforms())
+    num_classes = len(full_dataset.classes)
+
+    with open(CLASS_MAP_PATH, "w") as f:
+        json.dump(full_dataset.class_to_idx, f, separators=(",", ":"))
+
+    train_size = int(0.95 * len(full_dataset))
+    train_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, len(full_dataset)-train_size])
+    
+    train_loader = DataLoader(
+        train_dataset, batch_size=BATCH_SIZE, shuffle=True, 
+        num_workers=os.cpu_count(), pin_memory=True, prefetch_factor=2, persistent_workers=True
+    )
+    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, pin_memory=True)
+
+    model = CLIPItemEmbedder(num_classes).to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
+    scaler = GradScaler('cuda')
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
+    stopper = EarlyStopper(patience=PATIENCE, min_delta=MIN_DELTA)
+
+    print(f"Starting training for {num_classes} classes...")
+
+    for epoch in range(EPOCHS):
+        model.train()
+        loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{EPOCHS}")
+        
+        for images, labels in loop:
+            images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)
+            optimizer.zero_grad(set_to_none=True)
+
+            with autocast('cuda'):
+                img_emb, _, scale = model(images, labels)
+                img_emb = F.normalize(img_emb, p=2, dim=-1)
+
+                all_ids = torch.arange(num_classes, device=device)
+                prototypes = F.normalize(model.item_prototypes(all_ids), p=2, dim=-1)
+
+                logits = scale * img_emb @ prototypes.t()
+                loss = F.cross_entropy(logits, labels)
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            loop.set_postfix(loss=f"{loss.item():.4f}")
+
+        scheduler.step()
+
+        model.eval()
+        correct, total = 0, 0
+        with torch.no_grad(), autocast('cuda'):
+            all_ids = torch.arange(num_classes).to(device)
+            prototypes = F.normalize(model.item_prototypes(all_ids), p=2, dim=-1)
+            
+            for images, labels in val_loader:
+                images, labels = images.to(device), labels.to(device)
+                img_emb, _, _ = model(images, labels)
+                preds = (img_emb @ prototypes.t()).argmax(dim=-1)
+                correct += (preds == labels).sum().item()
+                total += labels.size(0)
+        
+        val_acc = 100 * correct / total
+        print(f"Validation Accuracy: {val_acc:.2f}%")
+
+        stop_now, is_best = stopper.check(val_acc, model)
+        if is_best:
+            torch.save(stopper.best_state, CLIP_BEST_PATH)
+            print("Successfully saved new best model weights.")
+        
+        if stop_now:
+            print(f"Stopping early. Best Accuracy was {stopper.best_accuracy:.2f}%")
+            break
+
+    print("Training finished.")
+
+if __name__ == "__main__":
+    train_clip()