Upload hierarchy_model.py with huggingface_hub

hierarchy_model.py

@@ -10,7 +10,7 @@ from sklearn.model_selection import train_test_split
 import re
 import requests
 from io import BytesIO
-from config import hierarchy_model_path, device, hierarchy_emb_dim, color_emb_dim, local_dataset_path, main_model_path, color_model_path
+import config
 
 # -------------------------
 # 1) Dataset
@@ -41,11 +41,11 @@ class HierarchyDataset(Dataset):
         
         # Check local image availability
         if use_local_images:
-            if 'local_image_path' not in dataframe.columns:
-                print("⚠️  Column 'local_image_path' not found. Using URLs.")
+            if config.column_local_image_path not in dataframe.columns:
+                print(f"⚠️  Column {config.column_local_image_path} not found. Using URLs.")
                 self.use_local_images = False
             else:
-                local_available = dataframe['local_image_path'].notna().sum()
+                local_available = dataframe[config.column_local_image_path].notna().sum()
                 total = len(dataframe)
                 print(f"📁 Local images available: {local_available}/{total} ({local_available/total*100:.1f}%)")
         
@@ -61,29 +61,24 @@ class HierarchyDataset(Dataset):
         row = self.dataframe.iloc[idx]
         
         # Try to load local image first
-        if self.use_local_images and pd.notna(row.get('local_image_path', '')):
-            local_path = row['local_image_path']
+        if self.use_local_images and pd.notna(row.get(config.column_local_image_path, '')):
+            local_path = row[config.column_local_image_path]
             image = Image.open(local_path).convert("RGB")
         # Check if image is a dictionary of bytes
-        elif isinstance(row['image_url'], dict):
-            image = Image.open(BytesIO(row['image_url']['bytes'])).convert('RGB')
+        elif isinstance(row[config.column_url_image], dict):
+            image = Image.open(BytesIO(row[config.column_url_image]['bytes'])).convert('RGB')
         # Otherwise, try to download from URL
         else:
-            try:
-                image = self._download_image(row['image_url'])
-            except Exception as e:
-                print(f"⚠️ Failed to load image {idx}: {e}")
-                # Create a blank image as fallback
-                image = Image.new('RGB', (224, 224), color='gray')
-        
+            image = self._download_image(row[config.column_url_image])
+
         # Apply transforms
         if hasattr(self, 'training_mode') and not self.training_mode:
             image = self.val_transform(image)
         else:
             image = self.transform(image)
 
-        description = row['text']  
-        hierarchy = row['hierarchy']
+        description = row[config.text_column]  
+        hierarchy = row[config.hierarchy_column]
 
         return image, description, hierarchy
     
@@ -190,7 +185,7 @@ class HierarchyExtractor:
             return self.class_to_idx[hierarchy]
         return None
     
-    def get_hierarchy_embedding(self, text, embed_dim=64):
+    def get_hierarchy_embedding(self, text, embed_dim=config.hierarchy_emb_dim):
         """Create embedding from hierarchy index"""
         hierarchy_idx = self.extract_hierarchy_idx(text)
         if hierarchy_idx is not None:
@@ -351,6 +346,9 @@ class Model(nn.Module):
         """Get text embeddings for a given text string or list of strings"""
         
         with torch.no_grad():
+            # Get the device of the model
+            model_device = next(self.parameters()).device
+            
             # Handle case where text is a list/tuple of hierarchies
             if isinstance(text, (list, tuple)):
                 # Process multiple hierarchies
@@ -365,7 +363,7 @@ class Model(nn.Module):
                         raise ValueError(f"Expected string, got {type(hierarchy_text)}: {hierarchy_text}")
                 
                 # Convert to tensor and move to device
-                hierarchy_indices = torch.tensor(hierarchy_indices, device=device)
+                hierarchy_indices = torch.tensor(hierarchy_indices, device=model_device)
                 
                 # Get text embeddings for all hierarchies
                 output = self.forward(hierarchy_indices=hierarchy_indices)
@@ -379,7 +377,7 @@ class Model(nn.Module):
                     raise ValueError(f"Could not extract hierarchy for text: '{text}'. Available classes: {self.hierarchy_extractor.hierarchy_classes}")
                 
                 # Convert to tensor and move to device
-                hierarchy_indices = torch.tensor([hierarchy_idx], device=device)
+                hierarchy_indices = torch.tensor([hierarchy_idx], device=model_device)
                 
                 # Get text embeddings
                 output = self.forward(hierarchy_indices=hierarchy_indices)
@@ -490,7 +488,7 @@ def collate_fn(batch, hierarchy_extractor):
     return {
         'image': images, 
         'hierarchy_indices': hierarchy_indices, 
-        'hierarchy': hierarchies
+        config.hierarchy_column: hierarchies
     }
 
 def calculate_accuracy(logits, target_hierarchies, hierarchy_classes):
@@ -525,7 +523,7 @@ def train_one_epoch(model, dataloader, optimizer, device, hierarchy_classes, sch
     for batch in pbar:
         images = batch['image'].to(device)
         hierarchy_indices = batch['hierarchy_indices'].to(device)
-        target_hierarchies = batch['hierarchy']
+        target_hierarchies = batch[config.hierarchy_column]
 
         # Set dataset to training mode
         if hasattr(dataloader.dataset, 'set_training_mode'):
@@ -590,7 +588,7 @@ def validate(model, dataloader, device, hierarchy_classes):
         for batch in pbar:
             images = batch['image'].to(device)
             hierarchy_indices = batch['hierarchy_indices'].to(device)
-            target_hierarchies = batch['hierarchy']
+            target_hierarchies = batch[config.hierarchy_column]
             
             # Set dataset to validation mode
             if hasattr(dataloader.dataset, 'set_training_mode'):
@@ -631,28 +629,24 @@ def validate(model, dataloader, device, hierarchy_classes):
 
 if __name__ == "__main__":
     # Configuration
-    CSV = "data/data_hierarchy_with_local_paths.csv"  
-    DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
-    BATCH = 16  
-    LR = 5e-5   
-    EPOCHS = 20 
-    VAL_SPLIT = 0.2 
-    EMB_DIM = 16
-    DROPOUT = 0.4
-    WEIGHT_DECAY = 1e-3
-
-    print(f"🚀 Starting hierarchical training on device: {DEVICE}")
-    print(f"📊 Config: {EPOCHS} epochs, batch={BATCH}, lr={LR}, embed_dim={EMB_DIM}")
+    device = config.device
+    batch_size = 16  
+    lr = 5e-5   
+    epochs = 20 
+    val_split = 0.2 
+    dropout = 0.4
+    weight_decay = 1e-3
+
+    print(f"🚀 Starting hierarchical training on device: {device}")
+    print(f"📊 Config: {epochs} epochs, batch={batch_size}, lr={lr}, embed_dim={config.hierarchy_emb_dim}")
 
     # Load dataset
-    print(f"📁 Using dataset: {CSV}")
-    df = pd.read_csv(CSV)
-    df = df[df['hierarchy'] != 'vest']
-    
+    print(f"📁 Using dataset: { config.local_dataset_path}")
+    df = pd.read_csv(config.local_dataset_path)    
     print(f"📁 Loaded {len(df)} samples")
     
     # Get unique hierarchy classes
-    hierarchy_classes = sorted(df['hierarchy'].unique().tolist())
+    hierarchy_classes = sorted(df[config.hierarchy_column].unique().tolist())
     print(f"📋 Found {len(hierarchy_classes)} hierarchy classes")
     
     # Create hierarchy extractor
@@ -661,9 +655,9 @@ if __name__ == "__main__":
     # Train/validation split
     train_df, val_df = train_test_split(
         df, 
-        test_size=VAL_SPLIT, 
+        test_size=val_split, 
         random_state=42,
-        stratify=df['hierarchy'] 
+        stratify=df[config.hierarchy_column] 
     )
     train_df = train_df.reset_index(drop=True)
     val_df = val_df.reset_index(drop=True)
@@ -677,13 +671,13 @@ if __name__ == "__main__":
     # Create data loaders
     train_dl = DataLoader(
         train_ds, 
-        batch_size=BATCH, 
+        batch_size=batch_size, 
         shuffle=True, 
         collate_fn=lambda batch: collate_fn(batch, hierarchy_extractor)
     )
     val_dl = DataLoader(
         val_ds, 
-        batch_size=BATCH, 
+        batch_size=batch_size, 
         shuffle=False, 
         collate_fn=lambda batch: collate_fn(batch, hierarchy_extractor)
     )
@@ -691,13 +685,13 @@ if __name__ == "__main__":
     # Create model
     model = Model(
         num_hierarchy_classes=len(hierarchy_classes),
-        embed_dim=EMB_DIM,
-        dropout=DROPOUT
-    ).to(DEVICE)
+        embed_dim=config.hierarchy_emb_dim,
+        dropout=dropout
+    ).to(device)
     
     # Optimizer and scheduler
-    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2, eta_min=LR/10)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2, eta_min=lr/10)
     
     print(f"🎯 Model parameters: {sum(p.numel() for p in model.parameters()):,}")
     print("\n" + "="*80)
@@ -706,15 +700,15 @@ if __name__ == "__main__":
     best_val_loss = float('inf')
     training_history = {'train_loss': [], 'val_loss': [], 'val_acc_img': [], 'val_acc_txt': []}
     
-    for e in range(EPOCHS):
-        print(f"\n🔄 Epoch {e+1}/{EPOCHS}")
+    for e in range(epochs):
+        print(f"\n🔄 Epoch {e+1}/{epochs}")
         print("-" * 50)
         
         # Training
-        train_metrics = train_one_epoch(model, train_dl, optimizer, DEVICE, hierarchy_classes, scheduler)
+        train_metrics = train_one_epoch(model, train_dl, optimizer, device, hierarchy_classes, scheduler)
         
         # Validation
-        val_metrics = validate(model, val_dl, DEVICE, hierarchy_classes)
+        val_metrics = validate(model, val_dl, device, hierarchy_classes)
         
         # Track history
         training_history['train_loss'].append(train_metrics['loss'])
@@ -741,10 +735,10 @@ if __name__ == "__main__":
                 'hierarchy_classes': hierarchy_classes,
                 'epoch': e+1,
                 'config': {
-                    'embed_dim': EMB_DIM,
-                    'dropout': DROPOUT
+                    'embed_dim': config.hierarchy_emb_dim,
+                    'dropout': dropout
                 }
-            }, "final_model_16.pth")
+            }, config.hierarchy_model_path)
         
         # Save model every 2 epochs
         if (e + 1) % 2 == 0:
@@ -755,10 +749,10 @@ if __name__ == "__main__":
                 'hierarchy_classes': hierarchy_classes,
                 'epoch': e+1,
                 'config': {
-                    'embed_dim': EMB_DIM,
-                    'dropout': DROPOUT
+                    'embed_dim': config.hierarchy_emb_dim,
+                    'dropout': dropout
                 }
-            }, f"model_checkpoint_epoch_{e+1}_16.pth")
+            }, f"model_checkpoint_epoch_{e+1}.pth")
 
     print("\n" + "="*80)
     print("🎉 Training completed!")