Add files copied from SPACES

Dockerfile

@@ -0,0 +1,45 @@
+FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    DEBIAN_FRONTEND=noninteractive
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 \
+    python3-pip \
+    python3-setuptools \
+    python3-dev \
+    build-essential \
+    git \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create a working directory
+WORKDIR /app
+
+# Installa dipendenze
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt 
+
+# Prepara directory di output con permessi corretti
+RUN mkdir -p  ./model_checkpoints  ./model    /tmp/huggingface_cache /app/embedding && \
+    chmod 777  ./model_checkpoints  ./model  /tmp/huggingface_cache  /app/embedding
+ 
+# Copia tutti i file necessari
+COPY *.py .
+COPY users.json .
+COPY run.sh .
+
+RUN chmod +x run.sh
+
+# Crea uno script per eseguire il processo e poi copiare i file
+#python generate_embeddings.py\n\
+# Esegui lo script
+
+# Esponi la porta per l'interfaccia web
+EXPOSE 7860
+
+# Set the entrypoint
+ENTRYPOINT ["./run.sh"]

README.md

@@ -0,0 +1,120 @@
+---
+title: Test01
+emoji: 🐨
+colorFrom: gray
+colorTo: red
+sdk: docker
+app_port: 7860
+startup_duration_timeout: 60m
+pinned: false
+license: apache-2.0
+---
+# User Embedding Model
+
+This repository contains a PyTorch model for generating user embeddings based on DMP (Data Management Platform) data. The model creates dense vector representations of users that can be used for recommendation systems, user clustering, and similarity searches.
+
+## Quick Start with Docker
+
+To run the model using Docker:
+
+```bash
+docker build -t user-embedding-model .
+
+docker run -v /path/to/your/data:/app/data \
+  -e DATA_PATH=/app/data/users.json \
+  -e NUM_EPOCHS=10 \
+  -e BATCH_SIZE=32 \
+  -v /path/to/output:/app/embeddings_output \
+  user-embedding-model
+```
+
+## Pushing to Hugging Face
+
+To automatically push the model to Hugging Face, add your credentials:
+
+```bash
+docker run -v /path/to/your/data:/app/data \
+  -e DATA_PATH=/app/data/users.json \
+  -e HF_REPO_ID="your-username/your-model-name" \
+  -e HF_TOKEN="your-huggingface-token" \
+  -v /path/to/output:/app/embeddings_output \
+  user-embedding-model
+```
+
+## Input Data Format
+
+The model expects user data in JSON format, with each user having DMP fields like:
+
+```json
+{
+  "dmp": {
+    "city": "milano",
+    "domains": ["example.com"],
+    "brands": ["brand1", "brand2"],
+    "clusters": ["cluster1", "cluster2"],
+    "industries": ["industry1"],
+    "tags": ["tag1", "tag2"],
+    "channels": ["channel1"],
+    "~click__host": "host1",
+    "~click__domain": "domain1",
+    "": {
+      "id": "user123"
+    }
+  }
+}
+```
+
+## Environment Variables
+
+- `DATA_PATH`: Path to your input JSON file (default: "users.json")
+- `NUM_EPOCHS`: Number of training epochs (default: 10)
+- `BATCH_SIZE`: Batch size for training (default: 32)
+- `LEARNING_RATE`: Learning rate for optimizer (default: 0.001)
+- `SAVE_INTERVAL`: Save checkpoint every N epochs (default: 2)
+- `HF_REPO_ID`: Hugging Face repository ID for uploading
+- `HF_TOKEN`: Hugging Face API token
+
+## Output
+
+The model generates:
+
+1. `embeddings.json`: User embeddings in JSON format
+2. `embeddings.npz`: User embeddings in NumPy format
+3. `vocabularies.json`: Vocabulary mappings
+4. `model.pth`: Trained PyTorch model
+5. `model_config.json`: Model configuration
+6. Hugging Face-compatible model files in the "huggingface" subdirectory
+
+## Hardware Requirements
+
+- **Recommended**: NVIDIA GPU with CUDA support
+- The code uses parallel processing for triplet generation to utilize all available CPU cores
+- For L40S GPU, recommended batch size: 32-64
+- Memory requirement: At least 8GB RAM
+
+## Model Architecture
+
+The model consists of:
+
+- Embedding layers for each user field
+- Sequential fully connected layers for dimensionality reduction
+- Output dimension: 256 (configurable)
+- Training method: Triplet margin loss using similar/dissimilar users
+
+## Performance Optimization
+
+The code includes several optimizations:
+
+- Parallel triplet generation using all available CPU cores
+- GPU acceleration for model training
+- Efficient memory handling for large datasets
+- TensorBoard integration for monitoring training
+
+## Troubleshooting
+
+If you encounter issues:
+
+1. Check that your input data follows the expected format
+2. Ensure you have sufficient memory for your dataset size
+3. For GPU errors, try reducing batch size
+4. Check the logs for detailed error messages

generate_model_backup.py

@@ -0,0 +1,1110 @@
+import os
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface_cache'
+#os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
+import multiprocessing
+try:
+    multiprocessing.set_start_method('spawn')
+except RuntimeError:
+    pass  # Il metodo è già stato impostato
+import json
+import torch
+import torch.nn as nn
+import numpy as np
+import pandas as pd
+from typing import List, Dict
+import logging
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import random
+from collections import defaultdict
+from pathlib import Path
+from tqdm import tqdm
+import os
+import multiprocessing
+from multiprocessing import Pool
+import psutil
+import argparse
+
+# =================================
+# CONFIGURABLE PARAMETERS
+# =================================
+# Define default parameters that can be overridden via environment variables
+DEFAULT_NUM_TRIPLETS = 150000      # Number of triplet examples to generate
+DEFAULT_NUM_EPOCHS = 20           # Number of training epochs
+DEFAULT_BATCH_SIZE = 64           # Batch size for training
+DEFAULT_LEARNING_RATE = 0.001     # Learning rate for optimizer
+DEFAULT_OUTPUT_DIM = 256          # Output dimension of embeddings
+DEFAULT_MAX_SEQ_LENGTH = 15       # Maximum sequence length
+DEFAULT_SAVE_INTERVAL = 2         # Save checkpoint every N epochs
+DEFAULT_DATA_PATH = "./users.json"  # Path to user data
+DEFAULT_OUTPUT_DIR = "./model"  # Output directory
+
+# Read parameters from environment variables
+NUM_TRIPLETS = int(os.environ.get("NUM_TRIPLETS", DEFAULT_NUM_TRIPLETS))
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", DEFAULT_NUM_EPOCHS))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", DEFAULT_BATCH_SIZE))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", DEFAULT_LEARNING_RATE))
+OUTPUT_DIM = int(os.environ.get("OUTPUT_DIM", DEFAULT_OUTPUT_DIM))
+MAX_SEQ_LENGTH = int(os.environ.get("MAX_SEQ_LENGTH", DEFAULT_MAX_SEQ_LENGTH))
+SAVE_INTERVAL = int(os.environ.get("SAVE_INTERVAL", DEFAULT_SAVE_INTERVAL))
+DATA_PATH = os.environ.get("DATA_PATH", DEFAULT_DATA_PATH)
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", DEFAULT_OUTPUT_DIR)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Get CUDA device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logging.info(f"Using device: {device}")
+
+# =================================
+# MODEL ARCHITECTURE
+# =================================
+class UserEmbeddingModel(nn.Module):
+    def __init__(self, vocab_sizes: Dict[str, int], embedding_dims: Dict[str, int], 
+                 output_dim: int = 256, max_sequence_length: int = 15,
+                 padded_fields_length: int = 10):
+        super().__init__()
+        
+        self.max_sequence_length = max_sequence_length
+        self.padded_fields_length = padded_fields_length
+        self.padded_fields = {'dmp_channels', 'dmp_tags', 'dmp_clusters'}
+        self.embedding_layers = nn.ModuleDict()
+        
+        # Create embedding layers for each field
+        for field, vocab_size in vocab_sizes.items():
+            self.embedding_layers[field] = nn.Embedding(
+                vocab_size, 
+                embedding_dims.get(field, 16),
+                padding_idx=0
+            )
+        
+        # Calculate total input dimension
+        self.total_input_dim = 0
+        for field, dim in embedding_dims.items():
+            if field in self.padded_fields:
+                self.total_input_dim += dim  # Single dimension for padded field
+            else:
+                self.total_input_dim += dim
+        
+        print(f"Total input dimension: {self.total_input_dim}")
+        
+        self.fc = nn.Sequential(
+            nn.Linear(self.total_input_dim, self.total_input_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(self.total_input_dim // 2, output_dim),
+            nn.LayerNorm(output_dim)
+        )
+
+    def _process_sequence(self, embedding_layer: nn.Embedding, indices: torch.Tensor, 
+                         field_name: str) -> torch.Tensor:
+        """Process normal sequences"""
+        batch_size = indices.size(0)
+        if indices.numel() == 0:
+            return torch.zeros(batch_size, embedding_layer.embedding_dim, device=indices.device)
+            
+        if field_name in ['dmp_city', 'dmp_domains']:
+            if indices.dim() == 1:
+                indices = indices.unsqueeze(0)
+            if indices.size(1) > 0:
+                return embedding_layer(indices[:, 0])
+            return torch.zeros(batch_size, embedding_layer.embedding_dim, device=indices.device)
+        
+        # Handle multiple sequences
+        embeddings = embedding_layer(indices)
+        return embeddings.mean(dim=1)  # [batch_size, emb_dim]
+
+    def _process_padded_sequence(self, embedding_layer: nn.Embedding, 
+                               indices: torch.Tensor) -> torch.Tensor:
+        """Process sequences with padding"""
+        batch_size = indices.size(0)
+        emb_dim = embedding_layer.embedding_dim
+        
+        # Generate embeddings
+        embeddings = embedding_layer(indices)  # [batch_size, seq_len, emb_dim]
+        
+        # Average along sequence dimension
+        mask = (indices != 0).float().unsqueeze(-1)
+        masked_embeddings = embeddings * mask
+        sum_mask = mask.sum(dim=1).clamp(min=1.0)
+        
+        return (masked_embeddings.sum(dim=1) / sum_mask)  # [batch_size, emb_dim]
+
+    def forward(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor:
+        batch_embeddings = []
+        
+        for field in ['dmp_city', 'dmp_domains', 'dmp_brands', 
+                     'dmp_clusters', 'dmp_industries', 'dmp_tags', 'dmp_channels', 
+                     'link_host', 'link_path']:
+            if field in inputs and field in self.embedding_layers:
+                if field in self.padded_fields:
+                    emb = self._process_padded_sequence(
+                        self.embedding_layers[field],
+                        inputs[field]
+                    )
+                else:
+                    emb = self._process_sequence(
+                        self.embedding_layers[field],
+                        inputs[field],
+                        field
+                    )
+                batch_embeddings.append(emb)
+        
+        combined = torch.cat(batch_embeddings, dim=1)
+        return self.fc(combined)
+
+# =================================
+# EMBEDDING PIPELINE
+# =================================
+class UserEmbeddingPipeline:
+    def __init__(self, output_dim: int = 256, max_sequence_length: int = 15):
+        self.output_dim = output_dim
+        self.max_sequence_length = max_sequence_length
+        self.model = None
+        self.vocab_maps = {}
+        
+        self.fields = [
+            'dmp_city', 'dmp_domains', 'dmp_brands', 
+            'dmp_clusters', 'dmp_industries', 'dmp_tags', 'dmp_channels',
+            'link_host', 'link_path'
+        ]
+        
+        # Map of new JSON fields to old field names used in the model
+        self.field_mapping = {
+            'dmp_city': ('dmp', 'city'),
+            'dmp_domains': ('dmp', 'domains'),
+            'dmp_brands': ('dmp', 'brands'),
+            'dmp_clusters': ('dmp', 'clusters'),
+            'dmp_industries': ('dmp', 'industries'),
+            'dmp_tags': ('dmp', 'tags'),
+            'dmp_channels': ('dmp', 'channels'),
+            'link_host': ('dmp', '~click__host'),
+            'link_path': ('dmp', '~click__domain')
+        }
+        
+        self.embedding_dims = {
+            'dmp_city': 16,
+            'dmp_domains': 16,
+            'dmp_brands': 32,
+            'dmp_clusters': 32,
+            'dmp_industries': 32,
+            'dmp_tags': 64,
+            'dmp_channels': 32,
+            'link_host': 32,
+            'link_path': 32
+        }
+
+    def _clean_value(self, value):
+        if isinstance(value, float) and np.isnan(value):
+            return []
+        if isinstance(value, str):
+            return [value.lower().strip()]
+        if isinstance(value, list):
+            return [str(v).lower().strip() for v in value if v is not None and str(v).strip()]
+        return []
+
+    def _get_field_from_user(self, user, field):
+        """Extract field value from new JSON user format"""
+        mapping = self.field_mapping.get(field, (field,))
+        value = user
+        
+        # Navigate through nested structure
+        for key in mapping:
+            if isinstance(value, dict):
+                value = value.get(key, {})
+            else:
+                # If not a dictionary and we're not at the last element
+                # of the mapping, return an empty list
+                value = []
+                break
+        
+        # If we've reached the end and have a value that's not a list but should be,
+        # convert it to a list
+        if field in {'dmp_brands', 'dmp_channels', 'dmp_clusters', 'dmp_industries', 'dmp_tags', 'link_host', 'link_path'} and not isinstance(value, list):
+            # If it's a string or other single value, put it in a list
+            if value and not isinstance(value, dict):
+                value = [value]
+            else:
+                value = []
+        
+        return value
+
+    def build_vocabularies(self, users_data: List[Dict]) -> Dict[str, Dict[str, int]]:
+        field_values = {field: {'<PAD>'} for field in self.fields}
+        
+        # Extract the 'user' field from the JSON structure for each record
+        users = []
+        for data in users_data:
+            # Check if there's raw_json.user
+            if 'raw_json' in data and 'user' in data['raw_json']:
+                users.append(data['raw_json']['user'])
+            # Check if there's user
+            elif 'user' in data:
+                users.append(data['user'])
+            else:
+                users.append(data)  # Assume it's already a user
+        
+        for user in users:
+            for field in self.fields:
+                values = self._clean_value(self._get_field_from_user(user, field))
+                field_values[field].update(values)
+        
+        self.vocab_maps = {
+            field: {val: idx for idx, val in enumerate(sorted(values))}
+            for field, values in field_values.items()
+        }
+        
+        return self.vocab_maps
+
+    def _prepare_input(self, user: Dict) -> Dict[str, torch.Tensor]:
+        inputs = {}
+        
+        for field in self.fields:
+            values = self._clean_value(self._get_field_from_user(user, field))
+            vocab = self.vocab_maps[field]
+            indices = [vocab.get(val, 0) for val in values]
+            inputs[field] = torch.tensor(indices, dtype=torch.long)
+            
+        return inputs
+
+    def initialize_model(self) -> None:
+        vocab_sizes = {field: len(vocab) for field, vocab in self.vocab_maps.items()}
+        
+        self.model = UserEmbeddingModel(
+            vocab_sizes=vocab_sizes,
+            embedding_dims=self.embedding_dims,
+            output_dim=self.output_dim,
+            max_sequence_length=self.max_sequence_length
+        )
+        self.model.to(device)
+        self.model.eval()
+
+    def generate_embeddings(self, users_data: List[Dict], batch_size: int = 32) -> Dict[str, np.ndarray]:
+        """Generate embeddings for all users"""
+        embeddings = {}
+        self.model.eval()  # Make sure model is in eval mode
+        
+        # Extract the 'user' field from the JSON structure for each record
+        users = []
+        user_ids = []
+        
+        for data in users_data:
+            # Check if there's raw_json.user
+            if 'raw_json' in data and 'user' in data['raw_json']:
+                user = data['raw_json']['user']
+                users.append(user)
+                # Use user.dmp[''].id as identifier
+                if 'dmp' in user and '' in user['dmp'] and 'id' in user['dmp']['']:
+                    user_ids.append(str(user['dmp']['']['id']))
+                else:
+                    # Fallback to uid or id if dmp.id is not available
+                    user_ids.append(str(user.get('uid', user.get('id', None))))
+            # Check if there's user
+            elif 'user' in data:
+                user = data['user']
+                users.append(user)
+                # Use user.dmp[''].id as identifier
+                if 'dmp' in user and '' in user['dmp'] and 'id' in user['dmp']['']:
+                    user_ids.append(str(user['dmp']['']['id']))
+                else:
+                    # Fallback to uid or id if dmp.id is not available
+                    user_ids.append(str(user.get('uid', user.get('id', None))))
+            else:
+                users.append(data)  # Assume it's already a user
+                # Use user.dmp[''].id as identifier
+                if 'dmp' in data and '' in data['dmp'] and 'id' in data['dmp']['']:
+                    user_ids.append(str(data['dmp']['']['id']))
+                else:
+                    # Fallback to uid or id if dmp.id is not available
+                    user_ids.append(str(data.get('uid', data.get('id', None))))
+        
+        with torch.no_grad():
+            for i in tqdm(range(0, len(users), batch_size), desc="Generating embeddings"):
+                batch_users = users[i:i+batch_size]
+                batch_ids = user_ids[i:i+batch_size]
+                batch_inputs = []
+                valid_indices = []
+                
+                for j, user in enumerate(batch_users):
+                    if batch_ids[j] is not None:
+                        batch_inputs.append(self._prepare_input(user))
+                        valid_indices.append(j)
+                
+                if batch_inputs:
+                    # Use the same collate function as training for a single batch
+                    anchor_batch, _, _ = collate_batch([(inputs, inputs, inputs) for inputs in batch_inputs])
+                    
+                    # Move data to device
+                    anchor_batch = {k: v.to(device) for k, v in anchor_batch.items()}
+                    
+                    # Generate embeddings
+                    batch_embeddings = self.model(anchor_batch).cpu()
+                    
+                    # Save embeddings
+                    for j, idx in enumerate(valid_indices):
+                        if batch_ids[idx]:  # Verify that id is not None or empty
+                            embeddings[batch_ids[idx]] = batch_embeddings[j].numpy()
+        
+        return embeddings
+
+    def save_embeddings(self, embeddings: Dict[str, np.ndarray], output_dir: str) -> None:
+        """Save embeddings to file"""
+        output_dir = Path(output_dir)
+        output_dir.mkdir(exist_ok=True)
+        
+        # Save embeddings as JSON
+        json_path = output_dir / 'embeddings.json'
+        with open(json_path, 'w') as f:
+            json_embeddings = {user_id: emb.tolist() for user_id, emb in embeddings.items()}
+            json.dump(json_embeddings, f)
+        
+        # Save embeddings as NPY
+        npy_path = output_dir / 'embeddings.npz'
+        np.savez_compressed(npy_path,
+                          embeddings=np.stack(list(embeddings.values())),
+                          user_ids=np.array(list(embeddings.keys())))
+        
+        # Save vocabularies
+        vocab_path = output_dir / 'vocabularies.json'
+        with open(vocab_path, 'w') as f:
+            json.dump(self.vocab_maps, f)
+            
+        logging.info(f"\nEmbeddings saved in {output_dir}:")
+        logging.info(f"- Embeddings JSON: {json_path}")
+        logging.info(f"- Embeddings NPY: {npy_path}")
+        logging.info(f"- Vocabularies: {vocab_path}")
+        
+    def save_model(self, output_dir: str) -> None:
+        """Save model in PyTorch format (.pth)"""
+        output_dir = Path(output_dir)
+        output_dir.mkdir(exist_ok=True)
+        
+        # Save path
+        model_path = output_dir / 'model.pth'
+        
+        # Prepare dictionary with model state and metadata
+        checkpoint = {
+            'model_state_dict': self.model.state_dict(),
+            'vocab_maps': self.vocab_maps,
+            'embedding_dims': self.embedding_dims,
+            'output_dim': self.output_dim,
+            'max_sequence_length': self.max_sequence_length
+        }
+        
+        # Save model
+        torch.save(checkpoint, model_path)
+        
+        logging.info(f"Model saved to: {model_path}")
+        
+        # Also save a configuration file for reference
+        config_info = {
+            'model_type': 'UserEmbeddingModel',
+            'vocab_sizes': {field: len(vocab) for field, vocab in self.vocab_maps.items()},
+            'embedding_dims': self.embedding_dims,
+            'output_dim': self.output_dim,
+            'max_sequence_length': self.max_sequence_length,
+            'padded_fields': list(self.model.padded_fields),
+            'fields': self.fields
+        }
+        
+        config_path = output_dir / 'model_config.json'
+        with open(config_path, 'w') as f:
+            json.dump(config_info, f, indent=2)
+            
+        logging.info(f"Model configuration saved to: {config_path}")
+        
+        # Save model in HuggingFace format
+        hf_dir = output_dir / 'huggingface'
+        hf_dir.mkdir(exist_ok=True)
+        
+        # Save model in HF format
+        torch.save(self.model.state_dict(), hf_dir / 'pytorch_model.bin')
+        
+        # Save config
+        with open(hf_dir / 'config.json', 'w') as f:
+            json.dump(config_info, f, indent=2)
+            
+        logging.info(f"Model saved in HuggingFace format to: {hf_dir}")
+        
+    def load_model(self, model_path: str) -> None:
+        """Load a previously saved model"""
+        checkpoint = torch.load(model_path, map_location=device)
+        
+        # Reload vocabularies and dimensions if needed
+        self.vocab_maps = checkpoint.get('vocab_maps', self.vocab_maps)
+        self.embedding_dims = checkpoint.get('embedding_dims', self.embedding_dims)
+        self.output_dim = checkpoint.get('output_dim', self.output_dim)
+        self.max_sequence_length = checkpoint.get('max_sequence_length', self.max_sequence_length)
+        
+        # Initialize model if not already done
+        if self.model is None:
+            self.initialize_model()
+            
+        # Load model weights
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model.to(device)
+        self.model.eval()
+        
+        logging.info(f"Model loaded from: {model_path}")
+
+# =================================
+# SIMILARITY AND TRIPLET GENERATION
+# =================================
+def calculate_similarity(user1, user2, pipeline):
+    try:
+        channels1 = set(str(c) for c in pipeline._get_field_from_user(user1, 'dmp_channels') if c is not None)
+        channels2 = set(str(c) for c in pipeline._get_field_from_user(user2, 'dmp_channels') if c is not None)
+        clusters1 = set(str(c) for c in pipeline._get_field_from_user(user1, 'dmp_clusters') if c is not None)
+        clusters2 = set(str(c) for c in pipeline._get_field_from_user(user2, 'dmp_clusters') if c is not None)
+        
+        channel_sim = len(channels1 & channels2) / max(1, len(channels1 | channels2))
+        cluster_sim = len(clusters1 & clusters2) / max(1, len(clusters1 | clusters2))
+        
+        return 0.5 * channel_sim + 0.5 * cluster_sim
+    except Exception as e:
+        logging.error(f"Error calculating similarity: {str(e)}")
+        return 0.0
+
+def process_batch_triplets(args):
+    try:
+        batch_idx, users, channel_index, cluster_index, num_triplets, pipeline = args
+        batch_triplets = []
+        
+        # Forza l'uso della CPU per tutti i calcoli
+        with torch.no_grad():
+            # Imposta temporaneamente il dispositivo su CPU per il calcolo delle similarità
+            temp_device = torch.device("cpu")
+            
+            for _ in range(num_triplets):
+                anchor_idx = random.randint(0, len(users)-1)
+                anchor_user = users[anchor_idx]
+                
+                # Find candidates that share channels or clusters
+                candidates = set()
+                for channel in pipeline._get_field_from_user(anchor_user, 'dmp_channels'):
+                    candidates.update(channel_index.get(str(channel), []))
+                for cluster in pipeline._get_field_from_user(anchor_user, 'dmp_clusters'):
+                    candidates.update(cluster_index.get(str(cluster), []))
+                    
+                # Remove anchor
+                candidates.discard(anchor_idx)
+                
+                # Find positive (similar) user
+                if not candidates:
+                    positive_idx = random.randint(0, len(users)-1)
+                else:
+                    # Calculate similarities for candidates
+                    similarities = []
+                    for idx in candidates:
+                        # Calcolo della similarità senza CUDA
+                        sim = cpu_calculate_similarity(anchor_user, users[idx], pipeline)
+                        if sim > 0:
+                            similarities.append((idx, sim))
+                    
+                    if not similarities:
+                        positive_idx = random.randint(0, len(users)-1)
+                    else:
+                        # Sort by similarity
+                        similarities.sort(key=lambda x: x[1], reverse=True)
+                        # Return one of the top K most similar
+                        top_k = min(10, len(similarities))
+                        positive_idx = similarities[random.randint(0, top_k-1)][0]
+                
+                # Find negative (dissimilar) user
+                max_attempts = 50
+                negative_idx = None
+                
+                for _ in range(max_attempts):
+                    idx = random.randint(0, len(users)-1)
+                    if idx != anchor_idx and idx != positive_idx:
+                        # Calcolo della similarità senza CUDA
+                        if cpu_calculate_similarity(anchor_user, users[idx], pipeline) < 0.1:
+                            negative_idx = idx
+                            break
+                
+                if negative_idx is None:
+                    negative_idx = random.randint(0, len(users)-1)
+                    
+                batch_triplets.append((anchor_idx, positive_idx, negative_idx))
+                
+        return batch_triplets
+    except Exception as e:
+        logging.error(f"Error in batch triplet generation: {str(e)}")
+        return []
+
+# Versione CPU della funzione di calcolo similarità
+def cpu_calculate_similarity(user1, user2, pipeline):
+    try:
+        channels1 = set(str(c) for c in pipeline._get_field_from_user(user1, 'dmp_channels') if c is not None)
+        channels2 = set(str(c) for c in pipeline._get_field_from_user(user2, 'dmp_channels') if c is not None)
+        clusters1 = set(str(c) for c in pipeline._get_field_from_user(user1, 'dmp_clusters') if c is not None)
+        clusters2 = set(str(c) for c in pipeline._get_field_from_user(user2, 'dmp_clusters') if c is not None)
+        
+        channel_sim = len(channels1 & channels2) / max(1, len(channels1 | channels2))
+        cluster_sim = len(clusters1 & clusters2) / max(1, len(clusters1 | clusters2))
+        
+        return 0.5 * channel_sim + 0.5 * cluster_sim
+    except Exception as e:
+        logging.error(f"Error calculating similarity: {str(e)}")
+        return 0.0    
+
+# =================================
+# DATASET AND DATALOADER
+# =================================
+class UserSimilarityDataset(Dataset):
+    def __init__(self, pipeline, users_data, num_triplets=10, num_workers=None):
+        self.triplets = []
+        logging.info("Initializing UserSimilarityDataset...")
+        
+        # Extract the 'user' field from the JSON structure for each record
+        self.users = []
+        for data in users_data:
+            # Check if there's raw_json.user
+            if 'raw_json' in data and 'user' in data['raw_json']:
+                self.users.append(data['raw_json']['user'])
+            # Check if there's user
+            elif 'user' in data:
+                self.users.append(data['user'])
+            else:
+                self.users.append(data)  # Assume it's already a user
+        
+        self.pipeline = pipeline
+        self.num_triplets = num_triplets
+        
+        # Determine number of workers for parallel processing
+        if num_workers is None:
+            num_workers = max(1, min(8, os.cpu_count()))
+        self.num_workers = num_workers
+        
+        # Pre-process inputs for each user
+        self.preprocessed_inputs = {}
+        for idx, user in enumerate(self.users):
+            self.preprocessed_inputs[idx] = pipeline._prepare_input(user)
+        
+        logging.info("Creating indexes for channels and clusters...")
+        self.channel_index = defaultdict(list)
+        self.cluster_index = defaultdict(list)
+        
+        for idx, user in enumerate(self.users):
+            channels = pipeline._get_field_from_user(user, 'dmp_channels')
+            clusters = pipeline._get_field_from_user(user, 'dmp_clusters')
+            
+            if channels:
+                channels = [str(c) for c in channels if c is not None]
+            if clusters:
+                clusters = [str(c) for c in clusters if c is not None]
+                
+            for channel in channels:
+                self.channel_index[channel].append(idx)
+            for cluster in clusters:
+                self.cluster_index[cluster].append(idx)
+        
+        logging.info(f"Found {len(self.channel_index)} unique channels and {len(self.cluster_index)} unique clusters")
+        logging.info(f"Generating triplets using {self.num_workers} worker processes...")
+        
+        self.triplets = self._generate_triplets_gpu(num_triplets)
+        logging.info(f"Generated {len(self.triplets)} triplets")
+
+    # Verifica che questo metodo sia definito correttamente
+    def __len__(self):
+        return len(self.triplets)
+    
+    def __getitem__(self, idx):
+        if idx >= len(self.triplets):
+            raise IndexError(f"Index {idx} out of range for dataset with {len(self.triplets)} triplets")
+        
+        anchor_idx, positive_idx, negative_idx = self.triplets[idx]
+        return (
+            self.preprocessed_inputs[anchor_idx],
+            self.preprocessed_inputs[positive_idx],
+            self.preprocessed_inputs[negative_idx]
+        )
+
+
+
+
+    def _generate_triplets_gpu(self, num_triplets):
+        """Generate triplets using a more reliable approach with batch processing"""
+        logging.info("Generating triplets with batch approach...")
+        
+        triplets = []
+        batch_size = 10  # Numero di triplette da generare per batch
+        num_batches = (num_triplets + batch_size - 1) // batch_size
+        
+        progress_bar = tqdm(
+            range(num_batches), 
+            desc="Generating triplet batches",
+            bar_format='{l_bar}{bar:10}{r_bar}'
+        )
+        
+        for _ in progress_bar:
+            batch_triplets = []
+            
+            # Genera un batch di triplette
+            for i in range(batch_size):
+                if len(triplets) >= num_triplets:
+                    break
+                    
+                # Seleziona anchor casuale
+                anchor_idx = random.randint(0, len(self.users)-1)
+                anchor_user = self.users[anchor_idx]
+                
+                # Trova candidati che condividono channels o clusters
+                candidates = set()
+                for channel in self.pipeline._get_field_from_user(anchor_user, 'dmp_channels'):
+                    candidates.update(self.channel_index.get(str(channel), []))
+                for cluster in self.pipeline._get_field_from_user(anchor_user, 'dmp_clusters'):
+                    candidates.update(self.cluster_index.get(str(cluster), []))
+                    
+                # Rimuovi l'anchor dai candidati
+                candidates.discard(anchor_idx)
+                
+                # Trova esempio positivo
+                if candidates:
+                    similarities = []
+                    for idx in list(candidates)[:50]:  # Limita la ricerca ai primi 50 candidati
+                        sim = calculate_similarity(anchor_user, self.users[idx], self.pipeline)
+                        if sim > 0:
+                            similarities.append((idx, sim))
+                    
+                    if similarities:
+                        similarities.sort(key=lambda x: x[1], reverse=True)
+                        top_k = min(10, len(similarities))
+                        positive_idx = similarities[random.randint(0, top_k-1)][0]
+                    else:
+                        positive_idx = random.randint(0, len(self.users)-1)
+                else:
+                    positive_idx = random.randint(0, len(self.users)-1)
+                
+                # Trova esempio negativo
+                attempts = 0
+                negative_idx = None
+                
+                while attempts < 20 and negative_idx is None:  # Ridotto a 20 tentativi
+                    idx = random.randint(0, len(self.users)-1)
+                    if idx != anchor_idx and idx != positive_idx:
+                        sim = calculate_similarity(anchor_user, self.users[idx], self.pipeline)
+                        if sim < 0.1:
+                            negative_idx = idx
+                            break
+                    attempts += 1
+                
+                if negative_idx is None:
+                    negative_idx = random.randint(0, len(self.users)-1)
+                
+                batch_triplets.append((anchor_idx, positive_idx, negative_idx))
+            
+            triplets.extend(batch_triplets)
+            
+        return triplets[:num_triplets]  # Assicurati di restituire esattamente num_triplets
+
+def collate_batch(batch):
+    """Custom collate function to properly handle tensor dimensions"""
+    anchor_inputs, positive_inputs, negative_inputs = zip(*batch)
+    
+    def process_group_inputs(group_inputs):
+        processed = {}
+        for field in group_inputs[0].keys():
+            # Find maximum length for this field in the batch
+            max_len = max(inputs[field].size(0) for inputs in group_inputs)
+            
+            # Create padded tensors
+            padded = torch.stack([
+                torch.cat([
+                    inputs[field],
+                    torch.zeros(max_len - inputs[field].size(0), dtype=torch.long)
+                ]) if inputs[field].size(0) < max_len else inputs[field][:max_len]
+                for inputs in group_inputs
+            ])
+            
+            processed[field] = padded
+            
+        return processed
+    
+    # Process each group (anchor, positive, negative)
+    anchor_batch = process_group_inputs(anchor_inputs)
+    positive_batch = process_group_inputs(positive_inputs)
+    negative_batch = process_group_inputs(negative_inputs)
+    
+    return anchor_batch, positive_batch, negative_batch
+
+# =================================
+# TRAINING FUNCTION
+# =================================
+
+def train_user_embeddings(model, users_data, pipeline, num_epochs=10, batch_size=32, lr=0.001, save_dir=None, save_interval=2, num_triplets=150):
+    """Main training of the model with proper batch handling and incremental saving"""
+    model.train()
+    model.to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    
+    # Add learning rate scheduler
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer, 
+        step_size=2,  # Decay every 2 epochs
+        gamma=0.9     # Multiply by 0.9 (10% reduction)
+    )
+    
+    # Determine number of CPU cores to use
+    num_cpu_cores = max(1, min(32, os.cpu_count()))
+    logging.info(f"Using {num_cpu_cores} CPU cores for data processing")
+    
+    # Prepare dataset and dataloader with custom collate function
+    dataset = UserSimilarityDataset(
+        pipeline, 
+        users_data, 
+        num_triplets=num_triplets,
+        num_workers=num_cpu_cores
+    )
+    
+    dataloader = DataLoader(
+        dataset, 
+        batch_size=batch_size, 
+        shuffle=True,
+        collate_fn=collate_batch,
+        num_workers=0,  # For loading batches
+        pin_memory=True  # Speed up data transfer to GPU
+    )
+    
+    # Loss function
+    criterion = torch.nn.TripletMarginLoss(margin=1.0)
+    
+    # Progress bar for epochs
+    epoch_pbar = tqdm(
+    range(num_epochs), 
+    desc="Training Progress",
+    bar_format='{l_bar}{bar:10}{r_bar}'
+    )
+    # Set up tensorboard for logging
+    try:
+        from torch.utils.tensorboard import SummaryWriter
+        log_dir = Path(save_dir) / "logs" if save_dir else Path("./logs")
+        log_dir.mkdir(exist_ok=True, parents=True)
+        writer = SummaryWriter(log_dir=log_dir)
+        tensorboard_available = True
+    except ImportError:
+        logging.warning("TensorBoard not available, skipping logging")
+        tensorboard_available = False
+    
+    for epoch in epoch_pbar:
+        total_loss = 0
+        num_batches = 0
+        
+        # Progress bar for batches
+        
+        total_batches = len(dataloader)
+        update_freq = max(1, total_batches // 10)  # Update approximately every 10%
+        batch_pbar = tqdm(
+            dataloader, 
+            desc=f"Epoch {epoch+1}/{num_epochs}",
+            leave=False,
+            miniters=update_freq,  # Only update progress bar every update_freq iterations
+            bar_format='{l_bar}{bar:10}{r_bar}',  # Simplified bar format
+            disable=True  # Completely disable the inner progress bar
+        )
+        
+        # First, create a single progress indicator for the entire epoch instead of the batch_pbar
+        epoch_progress = tqdm(
+            total=len(dataloader),
+            desc=f"Epoch {epoch+1}/{num_epochs}",
+            leave=True,
+            bar_format='{l_bar}{bar:10}{r_bar}'
+        )
+
+        # Then use this updated batch processing loop
+        for batch_idx, batch_inputs in enumerate(dataloader):
+            try:
+                # Each element in the batch is already a dict of padded tensors
+                anchor_batch, positive_batch, negative_batch = batch_inputs
+                
+                # Move data to device (GPU if available)
+                anchor_batch = {k: v.to(device) for k, v in anchor_batch.items()}
+                positive_batch = {k: v.to(device) for k, v in positive_batch.items()}
+                negative_batch = {k: v.to(device) for k, v in negative_batch.items()}
+                
+                # Generate embeddings
+                anchor_emb = model(anchor_batch)
+                positive_emb = model(positive_batch)
+                negative_emb = model(negative_batch)
+                
+                # Calculate loss
+                loss = criterion(anchor_emb, positive_emb, negative_emb)
+                
+                # Backward and optimize
+                optimizer.zero_grad()
+                loss.backward()
+                # Add gradient clipping
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                optimizer.step()
+                
+                total_loss += loss.item()
+                num_batches += 1
+                
+                # Update the epoch progress bar only at 10% intervals or at the end
+                update_interval = max(1, len(dataloader) // 10)
+                if (batch_idx + 1) % update_interval == 0 or batch_idx == len(dataloader) - 1:
+                    # Update progress
+                    remaining = min(update_interval, len(dataloader) - epoch_progress.n)
+                    epoch_progress.update(remaining)
+                    # Update stats with current average loss
+                    current_avg_loss = total_loss / num_batches
+                    epoch_progress.set_postfix(avg_loss=f"{current_avg_loss:.4f}", 
+                                              last_batch_loss=f"{loss.item():.4f}")
+                    
+            except Exception as e:
+                logging.error(f"Error during batch processing: {str(e)}")
+                logging.error(f"Batch details: {str(e.__class__.__name__)}")
+                continue
+
+        # Close the progress bar at the end of the epoch
+        epoch_progress.close()  
+
+        avg_loss = total_loss / max(1, num_batches)
+        # Update epoch progress bar with average loss
+        epoch_pbar.set_postfix(avg_loss=f"{avg_loss:.4f}")
+        
+        # Log to tensorboard
+        if tensorboard_available:
+            writer.add_scalar('Loss/train', avg_loss, epoch)
+        
+        # Step the learning rate scheduler at the end of each epoch
+        scheduler.step()
+        
+        # Incremental model saving if requested
+        if save_dir and (epoch + 1) % save_interval == 0:
+            checkpoint = {
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'loss': avg_loss,
+                'scheduler_state_dict': scheduler.state_dict()  # Save scheduler state
+            }
+            
+            save_path = Path(save_dir) / f'model_checkpoint_epoch_{epoch+1}.pth'
+            torch.save(checkpoint, save_path)
+            logging.info(f"Checkpoint saved at epoch {epoch+1}: {save_path}")
+    
+    if tensorboard_available:
+        writer.close()
+    
+    return model
+    
+# =================================
+# MAIN FUNCTION
+# =================================
+def main():
+    # Configuration
+    output_dir = Path(OUTPUT_DIR)
+    
+    # Check if CUDA is available
+    cuda_available = torch.cuda.is_available()
+    logging.info(f"CUDA available: {cuda_available}")
+    if cuda_available:
+        logging.info(f"CUDA device: {torch.cuda.get_device_name(0)}")
+        logging.info(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+    
+    # CPU info
+    cpu_count = os.cpu_count()
+    memory_info = psutil.virtual_memory()
+    logging.info(f"CPU cores: {cpu_count}")
+    logging.info(f"System memory: {memory_info.total / 1e9:.2f} GB")
+    
+    # Print configuration
+    logging.info("Running with the following configuration:")
+    logging.info(f"- Number of triplets: {NUM_TRIPLETS}")
+    logging.info(f"- Number of epochs: {NUM_EPOCHS}")
+    logging.info(f"- Batch size: {BATCH_SIZE}")
+    logging.info(f"- Learning rate: {LEARNING_RATE}")
+    logging.info(f"- Output dimension: {OUTPUT_DIM}")
+    logging.info(f"- Data path: {DATA_PATH}")
+    logging.info(f"- Output directory: {OUTPUT_DIR}")
+    
+    # Load data
+    logging.info("Loading user data...")
+    try:
+        try:
+            # First try loading as normal JSON
+            with open(DATA_PATH, 'r') as f:
+                json_data = json.load(f)
+                
+            # Handle both cases: array of users or single object
+            if isinstance(json_data, list):
+                users_data = json_data
+            elif isinstance(json_data, dict):
+                # If it's a single record, put it in a list
+                users_data = [json_data]
+            else:
+                raise ValueError("Unsupported JSON format")
+                
+        except json.JSONDecodeError:
+            # If it fails, the file might be a non-standard JSON (objects separated by commas)
+            logging.info("Detected possible non-standard JSON format, attempting correction...")
+            with open(DATA_PATH, 'r') as f:
+                text = f.read().strip()
+                
+            # Add square brackets to create a valid JSON array
+            if not text.startswith('['):
+                text = '[' + text
+            if not text.endswith(']'):
+                text = text + ']'
+                
+            # Try loading the corrected text
+            users_data = json.loads(text)
+            logging.info("JSON format successfully corrected")
+            
+        logging.info(f"Loaded {len(users_data)} records")
+    except FileNotFoundError:
+        logging.error(f"File {DATA_PATH} not found!")
+        return
+    except Exception as e:
+        logging.error(f"Unable to load file: {str(e)}")
+        return
+
+    # Initialize pipeline
+    logging.info("Initializing pipeline...")
+    pipeline = UserEmbeddingPipeline(
+        output_dim=OUTPUT_DIM,
+        max_sequence_length=MAX_SEQ_LENGTH
+    )
+    
+    # Build vocabularies
+    logging.info("Building vocabularies...")
+    try:
+        pipeline.build_vocabularies(users_data)
+        vocab_sizes = {field: len(vocab) for field, vocab in pipeline.vocab_maps.items()}
+        logging.info(f"Vocabulary sizes: {vocab_sizes}")
+    except Exception as e:
+        logging.error(f"Error building vocabularies: {str(e)}")
+        return
+
+    # Initialize model
+    logging.info("Initializing model...")
+    try:
+        pipeline.initialize_model()
+        logging.info("Model initialized successfully")
+    except Exception as e:
+        logging.error(f"Error initializing model: {str(e)}")
+        return
+
+    # Training
+    logging.info("Starting training...")
+    try:
+        # Create directory for checkpoints if it doesn't exist
+        model_dir = output_dir / "model_checkpoints"
+        model_dir.mkdir(exist_ok=True, parents=True)
+        
+        model = train_user_embeddings(
+            pipeline.model,
+            users_data,
+            pipeline,  # Pass the pipeline to training
+            num_epochs=NUM_EPOCHS,
+            batch_size=BATCH_SIZE,
+            lr=LEARNING_RATE,
+            save_dir=model_dir,  # Add incremental saving
+            save_interval=SAVE_INTERVAL,  # Save every N epochs
+            num_triplets=NUM_TRIPLETS
+        )
+        logging.info("Training completed")
+        pipeline.model = model
+        
+        # Save only the model file
+        logging.info("Saving model...")
+        
+        # Create output directory
+        output_dir.mkdir(exist_ok=True)
+        
+        # Save path
+        model_path = output_dir / 'model.pth'
+        
+        # Prepare dictionary with model state and metadata
+        checkpoint = {
+            'model_state_dict': pipeline.model.state_dict(),
+            'vocab_maps': pipeline.vocab_maps,
+            'embedding_dims': pipeline.embedding_dims,
+            'output_dim': pipeline.output_dim,
+            'max_sequence_length': pipeline.max_sequence_length
+        }
+        
+        # Save model
+        torch.save(checkpoint, model_path)
+        
+        logging.info(f"Model saved to: {model_path}")
+        
+        # Also save a configuration file for reference
+        config_info = {
+            'model_type': 'UserEmbeddingModel',
+            'vocab_sizes': {field: len(vocab) for field, vocab in pipeline.vocab_maps.items()},
+            'embedding_dims': pipeline.embedding_dims,
+            'output_dim': pipeline.output_dim,
+            'max_sequence_length': pipeline.max_sequence_length,
+            'padded_fields': list(pipeline.model.padded_fields),
+            'fields': pipeline.fields
+        }
+        
+        config_path = output_dir / 'model_config.json'
+        with open(config_path, 'w') as f:
+            json.dump(config_info, f, indent=2)
+            
+        logging.info(f"Model configuration saved to: {config_path}")
+        
+        # Only save in HuggingFace format if requested
+        save_hf = os.environ.get("SAVE_HF_FORMAT", "false").lower() == "true"
+        if save_hf:
+            logging.info("Saving in HuggingFace format...")
+            # Save model in HuggingFace format
+            hf_dir = output_dir / 'huggingface'
+            hf_dir.mkdir(exist_ok=True)
+            
+            # Save model in HF format
+            torch.save(pipeline.model.state_dict(), hf_dir / 'pytorch_model.bin')
+            
+            # Save config
+            with open(hf_dir / 'config.json', 'w') as f:
+                json.dump(config_info, f, indent=2)
+                
+            logging.info(f"Model saved in HuggingFace format to: {hf_dir}")
+        
+        # Push to HuggingFace if environment variable is set
+        hf_repo_id = os.environ.get("HF_REPO_ID")
+        hf_token = os.environ.get("HF_TOKEN")
+        
+        if save_hf and hf_repo_id and hf_token:
+            try:
+                from huggingface_hub import HfApi
+                
+                logging.info(f"Pushing model to HuggingFace: {hf_repo_id}")
+                api = HfApi()
+                
+                # Push the model directory
+                api.create_repo(
+                    repo_id=hf_repo_id,
+                    token=hf_token,
+                    exist_ok=True,
+                    private=True
+                )
+                
+                # Upload the files
+                for file_path in (output_dir / "huggingface").glob("**/*"):
+                    if file_path.is_file():
+                        api.upload_file(
+                            path_or_fileobj=str(file_path),
+                            path_in_repo=file_path.relative_to(output_dir / "huggingface"),
+                            repo_id=hf_repo_id,
+                            token=hf_token
+                        )
+                
+                logging.info(f"Model successfully pushed to HuggingFace: {hf_repo_id}")
+            except Exception as e:
+                logging.error(f"Error pushing to HuggingFace: {str(e)}")
+    
+    except Exception as e:
+        logging.error(f"Error during training or saving: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return
+
+    logging.info("Process completed successfully!")
+
+if __name__ == "__main__":
+    main()

generate_model_gpu.py

@@ -0,0 +1,1265 @@
+import os
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface_cache'
+#os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
+import multiprocessing
+try:
+    multiprocessing.set_start_method('spawn')
+except RuntimeError:
+    pass  # Il metodo è già stato impostato
+import json
+import torch
+import torch.nn as nn
+import numpy as np
+import pandas as pd
+from typing import List, Dict
+import logging
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import random
+from collections import defaultdict
+from pathlib import Path
+from tqdm import tqdm
+import os
+import multiprocessing
+from multiprocessing import Pool
+import psutil
+import argparse
+
+# =================================
+# CONFIGURABLE PARAMETERS
+# =================================
+# Define default parameters that can be overridden via environment variables
+DEFAULT_NUM_TRIPLETS = 150      # Number of triplet examples to generate
+DEFAULT_NUM_EPOCHS = 1           # Number of training epochs
+DEFAULT_BATCH_SIZE = 64           # Batch size for training
+DEFAULT_LEARNING_RATE = 0.001     # Learning rate for optimizer
+DEFAULT_OUTPUT_DIM = 256          # Output dimension of embeddings
+DEFAULT_MAX_SEQ_LENGTH = 15       # Maximum sequence length
+DEFAULT_SAVE_INTERVAL = 2         # Save checkpoint every N epochs
+DEFAULT_DATA_PATH = "./users.json"  # Path to user data
+DEFAULT_OUTPUT_DIR = "./model"  # Output directory
+
+# Read parameters from environment variables
+NUM_TRIPLETS = int(os.environ.get("NUM_TRIPLETS", DEFAULT_NUM_TRIPLETS))
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", DEFAULT_NUM_EPOCHS))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", DEFAULT_BATCH_SIZE))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", DEFAULT_LEARNING_RATE))
+OUTPUT_DIM = int(os.environ.get("OUTPUT_DIM", DEFAULT_OUTPUT_DIM))
+MAX_SEQ_LENGTH = int(os.environ.get("MAX_SEQ_LENGTH", DEFAULT_MAX_SEQ_LENGTH))
+SAVE_INTERVAL = int(os.environ.get("SAVE_INTERVAL", DEFAULT_SAVE_INTERVAL))
+DATA_PATH = os.environ.get("DATA_PATH", DEFAULT_DATA_PATH)
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", DEFAULT_OUTPUT_DIR)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Get CUDA device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logging.info(f"Using device: {device}")
+
+# =================================
+# MODEL ARCHITECTURE
+# =================================
+class UserEmbeddingModel(nn.Module):
+    def __init__(self, vocab_sizes: Dict[str, int], embedding_dims: Dict[str, int], 
+                 output_dim: int = 256, max_sequence_length: int = 15,
+                 padded_fields_length: int = 10):
+        super().__init__()
+        
+        self.max_sequence_length = max_sequence_length
+        self.padded_fields_length = padded_fields_length
+        self.padded_fields = {'dmp_channels', 'dmp_tags', 'dmp_clusters'}
+        self.embedding_layers = nn.ModuleDict()
+        
+        # Create embedding layers for each field
+        for field, vocab_size in vocab_sizes.items():
+            self.embedding_layers[field] = nn.Embedding(
+                vocab_size, 
+                embedding_dims.get(field, 16),
+                padding_idx=0
+            )
+        
+        # Calculate total input dimension
+        self.total_input_dim = 0
+        for field, dim in embedding_dims.items():
+            if field in self.padded_fields:
+                self.total_input_dim += dim  # Single dimension for padded field
+            else:
+                self.total_input_dim += dim
+        
+        print(f"Total input dimension: {self.total_input_dim}")
+        
+        self.fc = nn.Sequential(
+            nn.Linear(self.total_input_dim, self.total_input_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(self.total_input_dim // 2, output_dim),
+            nn.LayerNorm(output_dim)
+        )
+
+    def _process_sequence(self, embedding_layer: nn.Embedding, indices: torch.Tensor, 
+                         field_name: str) -> torch.Tensor:
+        """Process normal sequences"""
+        batch_size = indices.size(0)
+        if indices.numel() == 0:
+            return torch.zeros(batch_size, embedding_layer.embedding_dim, device=indices.device)
+            
+        if field_name in ['dmp_city', 'dmp_domains']:
+            if indices.dim() == 1:
+                indices = indices.unsqueeze(0)
+            if indices.size(1) > 0:
+                return embedding_layer(indices[:, 0])
+            return torch.zeros(batch_size, embedding_layer.embedding_dim, device=indices.device)
+        
+        # Handle multiple sequences
+        embeddings = embedding_layer(indices)
+        return embeddings.mean(dim=1)  # [batch_size, emb_dim]
+
+    def _process_padded_sequence(self, embedding_layer: nn.Embedding, 
+                               indices: torch.Tensor) -> torch.Tensor:
+        """Process sequences with padding"""
+        batch_size = indices.size(0)
+        emb_dim = embedding_layer.embedding_dim
+        
+        # Generate embeddings
+        embeddings = embedding_layer(indices)  # [batch_size, seq_len, emb_dim]
+        
+        # Average along sequence dimension
+        mask = (indices != 0).float().unsqueeze(-1)
+        masked_embeddings = embeddings * mask
+        sum_mask = mask.sum(dim=1).clamp(min=1.0)
+        
+        return (masked_embeddings.sum(dim=1) / sum_mask)  # [batch_size, emb_dim]
+
+    def forward(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor:
+            batch_embeddings = []
+            
+            for field in ['dmp_city', 'source', 'dmp_brands',  # modificato: rimosso 'dmp_domains', aggiunto 'source'
+                         'dmp_clusters', 'dmp_industries', 'dmp_tags', 'dmp_channels', 
+                         'device']:   # aggiunto 'device'
+                if field in inputs and field in self.embedding_layers:
+                    if field in self.padded_fields:
+                        emb = self._process_padded_sequence(
+                            self.embedding_layers[field],
+                            inputs[field]
+                        )
+                    else:
+                        emb = self._process_sequence(
+                            self.embedding_layers[field],
+                            inputs[field],
+                            field
+                        )
+                    batch_embeddings.append(emb)
+            
+            combined = torch.cat(batch_embeddings, dim=1)
+            return self.fc(combined)
+
+# =================================
+# EMBEDDING PIPELINE
+# =================================
+class UserEmbeddingPipeline:
+    def __init__(self, output_dim: int = 256, max_sequence_length: int = 15):
+        self.output_dim = output_dim
+        self.max_sequence_length = max_sequence_length
+        self.model = None
+        self.vocab_maps = {}
+        
+       
+        self.fields = [
+            'dmp_city', 'source', 'dmp_brands',  # 'dmp_domains' rimosso, 'source' aggiunto
+            'dmp_clusters', 'dmp_industries', 'dmp_tags', 'dmp_channels',
+            'device'  # 'device' aggiunto
+        ]
+        
+        # Map of new JSON fields to old field names used in the model
+        self.field_mapping = {
+            'dmp_city': ('dmp', 'city'),
+            'source': ('dmp', '', 'source'),
+            'dmp_brands': ('dmp', 'brands'),
+            'dmp_clusters': ('dmp', 'clusters'),
+            'dmp_industries': ('dmp', 'industries'),
+            'dmp_tags': ('dmp', 'tags'),
+            'dmp_channels': ('dmp', 'channels'),
+            'device': ('device',)  # Nuovo campo device
+        }
+        
+        self.embedding_dims = {
+            'dmp_city': 8,
+            'source': 8,  # Dimensione per source
+            'dmp_brands': 32,
+            'dmp_clusters': 64,
+            'dmp_industries': 32,
+            'dmp_tags': 128,
+            'dmp_channels': 64,
+            'device': 8  # Dimensione per device
+        }
+
+    def _clean_value(self, value):
+        if isinstance(value, float) and np.isnan(value):
+            return []
+        if isinstance(value, str):
+            return [value.lower().strip()]
+        if isinstance(value, list):
+            return [str(v).lower().strip() for v in value if v is not None and str(v).strip()]
+        return []
+
+    def _get_field_from_user(self, user, field):
+        """Extract field value from new JSON user format"""
+        mapping = self.field_mapping.get(field, (field,))
+        value = user
+        
+        # Navigate through nested structure
+        for key in mapping:
+            if isinstance(value, dict):
+                value = value.get(key, {})
+            else:
+                # If not a dictionary and we're not at the last element
+                # of the mapping, return an empty list
+                value = []
+                break
+        
+        # If we've reached the end and have a value that's not a list but should be,
+        # convert it to a list
+        if field in {'dmp_brands', 'dmp_channels', 'dmp_clusters', 'dmp_industries', 'dmp_tags'} and not isinstance(value, list):
+            # If it's a string or other single value, put it in a list
+            if value and not isinstance(value, dict):
+                value = [value]
+            else:
+                value = []
+        
+        return value
+
+    def build_vocabularies(self, users_data: List[Dict]) -> Dict[str, Dict[str, int]]:
+        field_values = {field: {'<PAD>'} for field in self.fields}
+        
+        # Extract the 'user' field from the JSON structure for each record
+        users = []
+        for data in users_data:
+            # Check if there's raw_json.user
+            if 'raw_json' in data and 'user' in data['raw_json']:
+                users.append(data['raw_json']['user'])
+            # Check if there's user
+            elif 'user' in data:
+                users.append(data['user'])
+            else:
+                users.append(data)  # Assume it's already a user
+        
+        for user in users:
+            for field in self.fields:
+                values = self._clean_value(self._get_field_from_user(user, field))
+                field_values[field].update(values)
+        
+        self.vocab_maps = {
+            field: {val: idx for idx, val in enumerate(sorted(values))}
+            for field, values in field_values.items()
+        }
+        
+        return self.vocab_maps
+
+    def _prepare_input(self, user: Dict) -> Dict[str, torch.Tensor]:
+        inputs = {}
+        
+        for field in self.fields:
+            values = self._clean_value(self._get_field_from_user(user, field))
+            vocab = self.vocab_maps[field]
+            indices = [vocab.get(val, 0) for val in values]
+            inputs[field] = torch.tensor(indices, dtype=torch.long)
+            
+        return inputs
+
+    def initialize_model(self) -> None:
+        vocab_sizes = {field: len(vocab) for field, vocab in self.vocab_maps.items()}
+        
+        self.model = UserEmbeddingModel(
+            vocab_sizes=vocab_sizes,
+            embedding_dims=self.embedding_dims,
+            output_dim=self.output_dim,
+            max_sequence_length=self.max_sequence_length
+        )
+        self.model.to(device)
+        self.model.eval()
+
+    def generate_embeddings(self, users_data: List[Dict], batch_size: int = 32) -> Dict[str, np.ndarray]:
+        """Generate embeddings for all users"""
+        embeddings = {}
+        self.model.eval()  # Make sure model is in eval mode
+        
+        # Extract the 'user' field from the JSON structure for each record
+        users = []
+        user_ids = []
+        
+        for data in users_data:
+            # Check if there's raw_json.user
+            if 'raw_json' in data and 'user' in data['raw_json']:
+                user = data['raw_json']['user']
+                users.append(user)
+                # Use user.dmp[''].id as identifier
+                if 'dmp' in user and '' in user['dmp'] and 'id' in user['dmp']['']:
+                    user_ids.append(str(user['dmp']['']['id']))
+                else:
+                    # Fallback to uid or id if dmp.id is not available
+                    user_ids.append(str(user.get('uid', user.get('id', None))))
+            # Check if there's user
+            elif 'user' in data:
+                user = data['user']
+                users.append(user)
+                # Use user.dmp[''].id as identifier
+                if 'dmp' in user and '' in user['dmp'] and 'id' in user['dmp']['']:
+                    user_ids.append(str(user['dmp']['']['id']))
+                else:
+                    # Fallback to uid or id if dmp.id is not available
+                    user_ids.append(str(user.get('uid', user.get('id', None))))
+            else:
+                users.append(data)  # Assume it's already a user
+                # Use user.dmp[''].id as identifier
+                if 'dmp' in data and '' in data['dmp'] and 'id' in data['dmp']['']:
+                    user_ids.append(str(data['dmp']['']['id']))
+                else:
+                    # Fallback to uid or id if dmp.id is not available
+                    user_ids.append(str(data.get('uid', data.get('id', None))))
+        
+        with torch.no_grad():
+            for i in tqdm(range(0, len(users), batch_size), desc="Generating embeddings"):
+                batch_users = users[i:i+batch_size]
+                batch_ids = user_ids[i:i+batch_size]
+                batch_inputs = []
+                valid_indices = []
+                
+                for j, user in enumerate(batch_users):
+                    if batch_ids[j] is not None:
+                        batch_inputs.append(self._prepare_input(user))
+                        valid_indices.append(j)
+                
+                if batch_inputs:
+                    # Use the same collate function as training for a single batch
+                    anchor_batch, _, _ = collate_batch([(inputs, inputs, inputs) for inputs in batch_inputs])
+                    
+                    # Move data to device
+                    anchor_batch = {k: v.to(device) for k, v in anchor_batch.items()}
+                    
+                    # Generate embeddings
+                    batch_embeddings = self.model(anchor_batch).cpu()
+                    
+                    # Save embeddings
+                    for j, idx in enumerate(valid_indices):
+                        if batch_ids[idx]:  # Verify that id is not None or empty
+                            embeddings[batch_ids[idx]] = batch_embeddings[j].numpy()
+        
+        return embeddings
+
+    def save_embeddings(self, embeddings: Dict[str, np.ndarray], output_dir: str) -> None:
+        """Save embeddings to file"""
+        output_dir = Path(output_dir)
+        output_dir.mkdir(exist_ok=True)
+        
+        # Save embeddings as JSON
+        json_path = output_dir / 'embeddings.json'
+        with open(json_path, 'w') as f:
+            json_embeddings = {user_id: emb.tolist() for user_id, emb in embeddings.items()}
+            json.dump(json_embeddings, f)
+        
+        # Save embeddings as NPY
+        npy_path = output_dir / 'embeddings.npz'
+        np.savez_compressed(npy_path,
+                          embeddings=np.stack(list(embeddings.values())),
+                          user_ids=np.array(list(embeddings.keys())))
+        
+        # Save vocabularies
+        vocab_path = output_dir / 'vocabularies.json'
+        with open(vocab_path, 'w') as f:
+            json.dump(self.vocab_maps, f)
+            
+        logging.info(f"\nEmbeddings saved in {output_dir}:")
+        logging.info(f"- Embeddings JSON: {json_path}")
+        logging.info(f"- Embeddings NPY: {npy_path}")
+        logging.info(f"- Vocabularies: {vocab_path}")
+        
+    def save_model(self, output_dir: str) -> None:
+        """Save model in PyTorch format (.pth)"""
+        output_dir = Path(output_dir)
+        output_dir.mkdir(exist_ok=True)
+        
+        # Save path
+        model_path = output_dir / 'model.pth'
+        
+        # Prepare dictionary with model state and metadata
+        checkpoint = {
+            'model_state_dict': self.model.state_dict(),
+            'vocab_maps': self.vocab_maps,
+            'embedding_dims': self.embedding_dims,
+            'output_dim': self.output_dim,
+            'max_sequence_length': self.max_sequence_length
+        }
+        
+        # Save model
+        torch.save(checkpoint, model_path)
+        
+        logging.info(f"Model saved to: {model_path}")
+        
+        # Also save a configuration file for reference
+        config_info = {
+            'model_type': 'UserEmbeddingModel',
+            'vocab_sizes': {field: len(vocab) for field, vocab in self.vocab_maps.items()},
+            'embedding_dims': self.embedding_dims,
+            'output_dim': self.output_dim,
+            'max_sequence_length': self.max_sequence_length,
+            'padded_fields': list(self.model.padded_fields),
+            'fields': self.fields
+        }
+        
+        config_path = output_dir / 'model_config.json'
+        with open(config_path, 'w') as f:
+            json.dump(config_info, f, indent=2)
+            
+        logging.info(f"Model configuration saved to: {config_path}")
+        
+        # Save model in HuggingFace format
+        hf_dir = output_dir / 'huggingface'
+        hf_dir.mkdir(exist_ok=True)
+        
+        # Save model in HF format
+        torch.save(self.model.state_dict(), hf_dir / 'pytorch_model.bin')
+        
+        # Save config
+        with open(hf_dir / 'config.json', 'w') as f:
+            json.dump(config_info, f, indent=2)
+            
+        logging.info(f"Model saved in HuggingFace format to: {hf_dir}")
+        
+    def load_model(self, model_path: str) -> None:
+        """Load a previously saved model"""
+        checkpoint = torch.load(model_path, map_location=device)
+        
+        # Reload vocabularies and dimensions if needed
+        self.vocab_maps = checkpoint.get('vocab_maps', self.vocab_maps)
+        self.embedding_dims = checkpoint.get('embedding_dims', self.embedding_dims)
+        self.output_dim = checkpoint.get('output_dim', self.output_dim)
+        self.max_sequence_length = checkpoint.get('max_sequence_length', self.max_sequence_length)
+        
+        # Initialize model if not already done
+        if self.model is None:
+            self.initialize_model()
+            
+        # Load model weights
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model.to(device)
+        self.model.eval()
+        
+        logging.info(f"Model loaded from: {model_path}")
+
+# =================================
+# SIMILARITY AND TRIPLET GENERATION
+# =================================
+
+def calculate_similarity(user1, user2, pipeline, filtered_tags=None):
+    try:
+        # Estrai i campi originali: channels e clusters
+        channels1 = set(str(c) for c in pipeline._get_field_from_user(user1, 'dmp_channels') if c is not None)
+        channels2 = set(str(c) for c in pipeline._get_field_from_user(user2, 'dmp_channels') if c is not None)
+        clusters1 = set(str(c) for c in pipeline._get_field_from_user(user1, 'dmp_clusters') if c is not None)
+        clusters2 = set(str(c) for c in pipeline._get_field_from_user(user2, 'dmp_clusters') if c is not None)
+        
+        # RIMOSSO domains1/domains2
+        
+        # Estrai i tag e applica il filtro se necessario
+        tags1 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user1, 'dmp_tags') if c is not None)
+        tags2 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user2, 'dmp_tags') if c is not None)
+        
+        # NUOVI CAMPI: source, brands, device
+        source1 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user1, 'source') if c is not None)
+        source2 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user2, 'source') if c is not None)
+        
+        brands1 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user1, 'dmp_brands') if c is not None)
+        brands2 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user2, 'dmp_brands') if c is not None)
+        
+        device1 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user1, 'device') if c is not None)
+        device2 = set(str(c).lower().strip() for c in pipeline._get_field_from_user(user2, 'device') if c is not None)
+        
+        if filtered_tags is not None:
+            # Filtra i tag usando solo quelli presenti nel set di tag filtrati
+            tags1 = {tag for tag in tags1 if tag in filtered_tags}
+            tags2 = {tag for tag in tags2 if tag in filtered_tags}
+        
+        # Calcola le similarità Jaccard per ogni campo
+        channel_sim = len(channels1 & channels2) / max(1, len(channels1 | channels2))
+        cluster_sim = len(clusters1 & clusters2) / max(1, len(clusters1 | channels2))
+      
+        tag_sim = len(tags1 & tags2) / max(1, len(tags1 | tags2))
+        
+        # Nuove similarità
+        source_sim = len(source1 & source2) / max(1, len(source1 | source2))
+        brands_sim = len(brands1 & brands2) / max(1, len(brands1 | brands2))
+        device_sim = len(device1 & device2) / max(1, len(device1 | device2))
+        
+        # Calcola la similarità totale con i pesi specificati:
+        # 6 per clusters, 5 per channels, 3 per tags,
+        # 3 per source, 5 per brands, 2 per device
+        total_weight = 6 + 5 + 3 + 2 + 5 + 3  # Somma dei pesi = 24
+        weighted_sim = (
+            6 * cluster_sim + 
+            5 * channel_sim + 
+            3 * tag_sim + 
+            2 * source_sim +  # Nuovo: source con peso 0.3 (3/24 ≈ 0.125)
+            5 * brands_sim +  # Nuovo: brands con peso 0.5 (5/24 ≈ 0.208) 
+            3 * device_sim    # Nuovo: device con peso 0.2 (2/24 ≈ 0.083)
+        ) / total_weight
+        
+        return weighted_sim
+    except Exception as e:
+        logging.error(f"Error calculating similarity: {str(e)}")
+        return 0.0
+
+
+def process_batch_triplets(args):
+    try:
+        batch_idx, users, channel_index, cluster_index, num_triplets, pipeline = args
+        batch_triplets = []
+        
+        # Forza l'uso della CPU per tutti i calcoli
+        with torch.no_grad():
+            # Imposta temporaneamente il dispositivo su CPU per il calcolo delle similarità
+            temp_device = torch.device("cpu")
+            
+            for _ in range(num_triplets):
+                anchor_idx = random.randint(0, len(users)-1)
+                anchor_user = users[anchor_idx]
+                
+                # Find candidates that share channels or clusters
+                candidates = set()
+                for channel in pipeline._get_field_from_user(anchor_user, 'dmp_channels'):
+                    candidates.update(channel_index.get(str(channel), []))
+                for cluster in pipeline._get_field_from_user(anchor_user, 'dmp_clusters'):
+                    candidates.update(cluster_index.get(str(cluster), []))
+                    
+                # Remove anchor
+                candidates.discard(anchor_idx)
+                
+                # Find positive (similar) user
+                if not candidates:
+                    positive_idx = random.randint(0, len(users)-1)
+                else:
+                    # Calculate similarities for candidates
+                    similarities = []
+                    for idx in candidates:
+                        # Calcolo della similarità senza CUDA
+                        sim = calculate_similarity(anchor_user, users[idx], pipeline)
+                        if sim > 0:
+                            similarities.append((idx, sim))
+                    
+                    if not similarities:
+                        positive_idx = random.randint(0, len(users)-1)
+                    else:
+                        # Sort by similarity
+                        similarities.sort(key=lambda x: x[1], reverse=True)
+                        # Return one of the top K most similar
+                        top_k = min(10, len(similarities))
+                        positive_idx = similarities[random.randint(0, top_k-1)][0]
+                
+                # Find negative (dissimilar) user
+                max_attempts = 50
+                negative_idx = None
+                
+                for _ in range(max_attempts):
+                    idx = random.randint(0, len(users)-1)
+                    if idx != anchor_idx and idx != positive_idx:
+                        # Calcolo della similarità senza CUDA
+                        if calculate_similarity(anchor_user, users[idx], pipeline) < 0.1:
+                            negative_idx = idx
+                            break
+                
+                if negative_idx is None:
+                    negative_idx = random.randint(0, len(users)-1)
+                    
+                batch_triplets.append((anchor_idx, positive_idx, negative_idx))
+                
+        return batch_triplets
+    except Exception as e:
+        logging.error(f"Error in batch triplet generation: {str(e)}")
+        return []
+
+
+# =================================
+# DATASET AND DATALOADER
+# =================================
+
+class UserSimilarityDataset(Dataset):
+    def __init__(self, pipeline, users_data, num_triplets=10, num_workers=None, filtered_tags=None):
+        self.triplets = []
+        self.filtered_tags = filtered_tags
+        logging.info("Initializing UserSimilarityDataset...")
+        
+        # Extract the 'user' field from the JSON structure for each record
+        self.users = []
+        for data in users_data:
+            # Check if there's raw_json.user
+            if 'raw_json' in data and 'user' in data['raw_json']:
+                self.users.append(data['raw_json']['user'])
+            # Check if there's user
+            elif 'user' in data:
+                self.users.append(data['user'])
+            else:
+                self.users.append(data)  # Assume it's already a user
+        
+        self.pipeline = pipeline
+        self.num_triplets = num_triplets
+        
+        # Determine number of workers for parallel processing
+        if num_workers is None:
+            num_workers = max(1, min(8, os.cpu_count()))
+        self.num_workers = num_workers
+        
+        # Pre-processa gli input per ogni utente
+        self.preprocessed_inputs = {}
+        for idx, user in enumerate(self.users):
+            self.preprocessed_inputs[idx] = pipeline._prepare_input(user)
+        
+        logging.info("Creating indexes for channels, clusters, tags, brands, source, and device...")
+        self.channel_index = defaultdict(list)
+        self.cluster_index = defaultdict(list)
+        self.tag_index = defaultdict(list)
+        # Rimosso domain_index
+        self.brands_index = defaultdict(list)  # Nuovo indice per brands
+        self.source_index = defaultdict(list)  # Nuovo indice per source
+        self.device_index = defaultdict(list)  # Nuovo indice per device
+        
+        for idx, user in enumerate(self.users):
+            channels = pipeline._get_field_from_user(user, 'dmp_channels')
+            clusters = pipeline._get_field_from_user(user, 'dmp_clusters')
+            tags = pipeline._get_field_from_user(user, 'dmp_tags')
+            # Rimosso domains
+            brands = pipeline._get_field_from_user(user, 'dmp_brands')  # Aggiunto brands
+            source = pipeline._get_field_from_user(user, 'source')      # Aggiunto source
+            device = pipeline._get_field_from_user(user, 'device')      # Aggiunto device
+            
+            if channels:
+                channels = [str(c) for c in channels if c is not None]
+            if clusters:
+                clusters = [str(c) for c in clusters if c is not None]
+            if tags:
+                tags = [str(c) for c in tags if c is not None]
+                # Filtra i tag se è stato fornito un set di tag filtrati
+                if self.filtered_tags:
+                    tags = [tag for tag in tags if tag in self.filtered_tags]
+            # Rimosso codice per domains
+            if brands:
+                brands = [str(c) for c in brands if c is not None]
+
+
+            if source:
+                if isinstance(source, str):
+                    source = [source]  # Se è una stringa singola, mettila in una lista
+                else:
+                    source = [str(c) for c in source if c is not None]  # Altrimenti trattala come prima
+
+            if device:
+                if isinstance(device, str):
+                    device = [device]  # Se è una stringa singola, mettila in una lista
+                else:
+                    device = [str(c) for c in device if c is not None]  
+                
+            for channel in channels:
+                self.channel_index[channel].append(idx)
+            for cluster in clusters:
+                self.cluster_index[cluster].append(idx)
+            for tag in tags:
+                self.tag_index[tag].append(idx)
+            # Rimosso ciclo per domains
+            for brand in brands:
+                self.brands_index[brand].append(idx)
+            for s in source:
+                self.source_index[s].append(idx)
+            for d in device:
+                self.device_index[d].append(idx)
+        
+        logging.info(f"Found {len(self.channel_index)} unique channels, {len(self.cluster_index)} unique clusters, {len(self.tag_index)} unique tags")
+        logging.info(f"Found {len(self.brands_index)} unique brands, {len(self.source_index)} unique sources, and {len(self.device_index)} unique devices")
+        
+        logging.info(f"Generating triplets using {self.num_workers} worker processes...")
+        
+        self.triplets = self._generate_triplets_gpu(num_triplets)
+        logging.info(f"Generated {len(self.triplets)} triplets")
+
+    def __len__(self):
+        return len(self.triplets)
+    
+    def __getitem__(self, idx):
+        if idx >= len(self.triplets):
+            raise IndexError(f"Index {idx} out of range for dataset with {len(self.triplets)} triplets")
+        
+        anchor_idx, positive_idx, negative_idx = self.triplets[idx]
+        return (
+            self.preprocessed_inputs[anchor_idx],
+            self.preprocessed_inputs[positive_idx],
+            self.preprocessed_inputs[negative_idx]
+        )
+
+    def _generate_triplets_gpu(self, num_triplets):
+        """Generate triplets using a more reliable approach with batch processing"""
+        logging.info("Generating triplets with batch approach...")
+        
+        triplets = []
+        batch_size = 10  # Numero di triplette da generare per batch
+        num_batches = (num_triplets + batch_size - 1) // batch_size
+        
+        progress_bar = tqdm(
+            range(num_batches), 
+            desc="Generating triplet batches",
+            bar_format='{l_bar}{bar:10}{r_bar}'
+        )
+        
+        for _ in progress_bar:
+            batch_triplets = []
+            
+            # Genera un batch di triplette
+            for i in range(batch_size):
+                if len(triplets) >= num_triplets:
+                    break
+                    
+                # Seleziona anchor casuale
+                anchor_idx = random.randint(0, len(self.users)-1)
+                anchor_user = self.users[anchor_idx]
+                
+                # Trova candidati che condividono channels, clusters, tags o domains
+                # Trova candidati che condividono channels, clusters, tags, brands, source, device
+                candidates = set()
+                for channel in self.pipeline._get_field_from_user(anchor_user, 'dmp_channels'):
+                    if channel is not None:
+                        candidates.update(self.channel_index.get(str(channel), []))
+                for cluster in self.pipeline._get_field_from_user(anchor_user, 'dmp_clusters'):
+                    if cluster is not None:
+                        candidates.update(self.cluster_index.get(str(cluster), []))
+                for tag in self.pipeline._get_field_from_user(anchor_user, 'dmp_tags'):
+                    if tag is not None and (self.filtered_tags is None or str(tag) in self.filtered_tags):
+                        candidates.update(self.tag_index.get(str(tag), []))
+                # Rimosso il loop per domains
+                
+                # Nuovi loop per brands, source, device
+                for brand in self.pipeline._get_field_from_user(anchor_user, 'dmp_brands'):
+                    if brand is not None:
+                        candidates.update(self.brands_index.get(str(brand), []))
+                for source in self.pipeline._get_field_from_user(anchor_user, 'source'):
+                    if source is not None:
+                        candidates.update(self.source_index.get(str(source), []))
+                for device in self.pipeline._get_field_from_user(anchor_user, 'device'):
+                    if device is not None:
+                        candidates.update(self.device_index.get(str(device), []))
+                    
+                # Rimuovi l'anchor dai candidati
+                candidates.discard(anchor_idx)
+                
+                # Trova esempio positivo
+                if candidates:
+                    similarities = []
+                    for idx in list(candidates)[:50]:  # Limita la ricerca ai primi 50 candidati
+                        sim = calculate_similarity(anchor_user, self.users[idx], self.pipeline, self.filtered_tags)
+                        if sim > 0:
+                            similarities.append((idx, sim))
+                    
+                    if similarities:
+                        similarities.sort(key=lambda x: x[1], reverse=True)
+                        top_k = min(10, len(similarities))
+                        positive_idx = similarities[random.randint(0, top_k-1)][0]
+                    else:
+                        positive_idx = random.randint(0, len(self.users)-1)
+                else:
+                    positive_idx = random.randint(0, len(self.users)-1)
+                
+                # Trova esempio negativo
+                attempts = 0
+                negative_idx = None
+                
+                while attempts < 20 and negative_idx is None:  # Ridotto a 20 tentativi
+                    idx = random.randint(0, len(self.users)-1)
+                    if idx != anchor_idx and idx != positive_idx:
+                        sim = calculate_similarity(anchor_user, self.users[idx], self.pipeline, self.filtered_tags)
+                        if sim < 0.1:
+                            negative_idx = idx
+                            break
+                    attempts += 1
+                
+                if negative_idx is None:
+                    negative_idx = random.randint(0, len(self.users)-1)
+                
+                batch_triplets.append((anchor_idx, positive_idx, negative_idx))
+            
+            triplets.extend(batch_triplets)
+            
+        return triplets[:num_triplets]  # Assicurati di restituire esattamente num_triplets 
+
+
+def collate_batch(batch):
+    """Custom collate function to properly handle tensor dimensions"""
+    anchor_inputs, positive_inputs, negative_inputs = zip(*batch)
+    
+    def process_group_inputs(group_inputs):
+        processed = {}
+        for field in group_inputs[0].keys():
+            # Find maximum length for this field in the batch
+            max_len = max(inputs[field].size(0) for inputs in group_inputs)
+            
+            # Create padded tensors
+            padded = torch.stack([
+                torch.cat([
+                    inputs[field],
+                    torch.zeros(max_len - inputs[field].size(0), dtype=torch.long)
+                ]) if inputs[field].size(0) < max_len else inputs[field][:max_len]
+                for inputs in group_inputs
+            ])
+            
+            processed[field] = padded
+            
+        return processed
+    
+    # Process each group (anchor, positive, negative)
+    anchor_batch = process_group_inputs(anchor_inputs)
+    positive_batch = process_group_inputs(positive_inputs)
+    negative_batch = process_group_inputs(negative_inputs)
+    
+    return anchor_batch, positive_batch, negative_batch
+
+# =================================
+# TRAINING FUNCTION
+# =================================
+
+def train_user_embeddings(model, users_data, pipeline, num_epochs=10, batch_size=32, lr=0.001, save_dir=None, save_interval=2, num_triplets=150):
+    """Main training of the model with proper batch handling and incremental saving"""
+    model.train()
+    model.to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    
+    # Add learning rate scheduler
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer, 
+        step_size=2,  # Decay every 2 epochs
+        gamma=0.9     # Multiply by 0.9 (10% reduction)
+    )
+    
+    # Determine number of CPU cores to use
+    num_cpu_cores = max(1, min(32, os.cpu_count()))
+    logging.info(f"Using {num_cpu_cores} CPU cores for data processing")
+    
+    # Prepare dataset and dataloader with custom collate function
+    dataset = UserSimilarityDataset(
+        pipeline, 
+        users_data, 
+        num_triplets=num_triplets,
+        num_workers=num_cpu_cores
+    )
+    
+    dataloader = DataLoader(
+        dataset, 
+        batch_size=batch_size, 
+        shuffle=True,
+        collate_fn=collate_batch,
+        num_workers=0,  # For loading batches
+        pin_memory=True  # Speed up data transfer to GPU
+    )
+    
+    # Loss function
+    criterion = torch.nn.TripletMarginLoss(margin=1.0)
+    
+    # Progress bar for epochs
+    epoch_pbar = tqdm(
+    range(num_epochs), 
+    desc="Training Progress",
+    bar_format='{l_bar}{bar:10}{r_bar}'
+    )
+    # Set up tensorboard for logging
+    try:
+        from torch.utils.tensorboard import SummaryWriter
+        log_dir = Path(save_dir) / "logs" if save_dir else Path("./logs")
+        log_dir.mkdir(exist_ok=True, parents=True)
+        writer = SummaryWriter(log_dir=log_dir)
+        tensorboard_available = True
+    except ImportError:
+        logging.warning("TensorBoard not available, skipping logging")
+        tensorboard_available = False
+    
+    for epoch in epoch_pbar:
+        total_loss = 0
+        num_batches = 0
+        
+        # Progress bar for batches
+        
+        total_batches = len(dataloader)
+        update_freq = max(1, total_batches // 10)  # Update approximately every 10%
+        batch_pbar = tqdm(
+            dataloader, 
+            desc=f"Epoch {epoch+1}/{num_epochs}",
+            leave=False,
+            miniters=update_freq,  # Only update progress bar every update_freq iterations
+            bar_format='{l_bar}{bar:10}{r_bar}',  # Simplified bar format
+            disable=True  # Completely disable the inner progress bar
+        )
+        
+        # First, create a single progress indicator for the entire epoch instead of the batch_pbar
+        epoch_progress = tqdm(
+            total=len(dataloader),
+            desc=f"Epoch {epoch+1}/{num_epochs}",
+            leave=True,
+            bar_format='{l_bar}{bar:10}{r_bar}'
+        )
+
+        # Then use this updated batch processing loop
+        for batch_idx, batch_inputs in enumerate(dataloader):
+            try:
+                # Each element in the batch is already a dict of padded tensors
+                anchor_batch, positive_batch, negative_batch = batch_inputs
+                
+                # Move data to device (GPU if available)
+                anchor_batch = {k: v.to(device) for k, v in anchor_batch.items()}
+                positive_batch = {k: v.to(device) for k, v in positive_batch.items()}
+                negative_batch = {k: v.to(device) for k, v in negative_batch.items()}
+                
+                # Generate embeddings
+                anchor_emb = model(anchor_batch)
+                positive_emb = model(positive_batch)
+                negative_emb = model(negative_batch)
+                
+                # Calculate loss
+                loss = criterion(anchor_emb, positive_emb, negative_emb)
+                
+                # Backward and optimize
+                optimizer.zero_grad()
+                loss.backward()
+                # Add gradient clipping
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                optimizer.step()
+                
+                total_loss += loss.item()
+                num_batches += 1
+                
+                # Update the epoch progress bar only at 10% intervals or at the end
+                update_interval = max(1, len(dataloader) // 10)
+                if (batch_idx + 1) % update_interval == 0 or batch_idx == len(dataloader) - 1:
+                    # Update progress
+                    remaining = min(update_interval, len(dataloader) - epoch_progress.n)
+                    epoch_progress.update(remaining)
+                    # Update stats with current average loss
+                    current_avg_loss = total_loss / num_batches
+                    epoch_progress.set_postfix(avg_loss=f"{current_avg_loss:.4f}", 
+                                              last_batch_loss=f"{loss.item():.4f}")
+                    
+            except Exception as e:
+                logging.error(f"Error during batch processing: {str(e)}")
+                logging.error(f"Batch details: {str(e.__class__.__name__)}")
+                continue
+
+        # Close the progress bar at the end of the epoch
+        epoch_progress.close()  
+
+        avg_loss = total_loss / max(1, num_batches)
+        # Update epoch progress bar with average loss
+        epoch_pbar.set_postfix(avg_loss=f"{avg_loss:.4f}")
+        
+        # Log to tensorboard
+        if tensorboard_available:
+            writer.add_scalar('Loss/train', avg_loss, epoch)
+        
+        # Step the learning rate scheduler at the end of each epoch
+        scheduler.step()
+        
+        # Incremental model saving if requested
+        if save_dir and (epoch + 1) % save_interval == 0:
+            checkpoint = {
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'loss': avg_loss,
+                'scheduler_state_dict': scheduler.state_dict()  # Save scheduler state
+            }
+            
+            save_path = Path(save_dir) / f'model_checkpoint_epoch_{epoch+1}.pth'
+            torch.save(checkpoint, save_path)
+            logging.info(f"Checkpoint saved at epoch {epoch+1}: {save_path}")
+    
+    if tensorboard_available:
+        writer.close()
+    
+    return model
+    
+
+def compute_tag_frequencies(pipeline, users_data):
+    """
+    Calcola le frequenze dei tag nel dataset.
+    
+    Args:
+        pipeline: Pipeline di embedding
+        users_data: Lista di utenti
+        
+    Returns:
+        dict: Dizionario con tag come chiavi e frequenze come valori
+        int: Numero totale di tag processati
+    """
+    logging.info("Calcolando frequenze dei tag...")
+    tag_frequencies = {}
+    total_tags = 0
+    
+    # Estrai i tag da tutti gli utenti
+    for data in users_data:
+        # Estrai l'utente dalla struttura JSON
+        if 'raw_json' in data and 'user' in data['raw_json']:
+            user = data['raw_json']['user']
+        elif 'user' in data:
+            user = data['user']
+        else:
+            user = data  # Assume che sia già un utente
+        
+        # Estrai e conta i tag
+        tags = pipeline._get_field_from_user(user, 'dmp_tags')
+        for tag in tags:
+            if tag is not None:
+                tag_str = str(tag).lower().strip()
+                tag_frequencies[tag_str] = tag_frequencies.get(tag_str, 0) + 1
+                total_tags += 1
+    
+    logging.info(f"Trovati {len(tag_frequencies)} tag unici su {total_tags} occorrenze totali")
+    return tag_frequencies, total_tags
+
+# funzione per filtrare i tag in base alla frequenza o al percentile
+
+def filter_tags_by_criteria(tag_frequencies, min_frequency=100, percentile=None):
+    """
+    Filtra i tag in base a criteri di frequenza o percentile.
+    
+    Args:
+        tag_frequencies: Dizionario con tag e frequenze
+        min_frequency: Frequenza minima richiesta (default: 100)
+        percentile: Se specificato, mantiene solo i tag fino al percentile indicato
+        
+    Returns:
+        set: Set di tag che soddisfano i criteri
+    """
+    if percentile is not None:
+        # Ordina i tag per frequenza e mantieni solo fino al percentile specificato
+        sorted_tags = sorted(tag_frequencies.items(), key=lambda x: x[1], reverse=True)
+        cutoff_index = int(len(sorted_tags) * (percentile / 100.0))
+        filtered_tags = {tag for tag, _ in sorted_tags[:cutoff_index]}
+        
+        min_freq_in_set = sorted_tags[cutoff_index-1][1] if cutoff_index > 0 else 0
+        logging.info(f"Filtrati tag al {percentile}° percentile. Mantenuti {len(filtered_tags)} tag con frequenza >= {min_freq_in_set}")
+    else:
+        # Filtra solo in base alla frequenza minima
+        filtered_tags = {tag for tag, freq in tag_frequencies.items() if freq >= min_frequency}
+        logging.info(f"Filtrati tag con frequenza < {min_frequency}. Mantenuti {len(filtered_tags)} tag")
+    
+    return filtered_tags
+    
+# =================================
+# MAIN FUNCTION
+# =================================
+def main():
+    # Configuration
+    output_dir = Path(OUTPUT_DIR)
+    
+    # Check if CUDA is available
+    cuda_available = torch.cuda.is_available()
+    logging.info(f"CUDA available: {cuda_available}")
+    if cuda_available:
+        logging.info(f"CUDA device: {torch.cuda.get_device_name(0)}")
+        logging.info(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+    
+    # CPU info
+    cpu_count = os.cpu_count()
+    memory_info = psutil.virtual_memory()
+    logging.info(f"CPU cores: {cpu_count}")
+    logging.info(f"System memory: {memory_info.total / 1e9:.2f} GB")
+    
+    # Print configuration
+    logging.info("Running with the following configuration:")
+    logging.info(f"- Number of triplets: {NUM_TRIPLETS}")
+    logging.info(f"- Number of epochs: {NUM_EPOCHS}")
+    logging.info(f"- Batch size: {BATCH_SIZE}")
+    logging.info(f"- Learning rate: {LEARNING_RATE}")
+    logging.info(f"- Output dimension: {OUTPUT_DIM}")
+    logging.info(f"- Data path: {DATA_PATH}")
+    logging.info(f"- Output directory: {OUTPUT_DIR}")
+    
+    # Load data
+    logging.info("Loading user data...")
+    try:
+        try:
+            # First try loading as normal JSON
+            with open(DATA_PATH, 'r') as f:
+                json_data = json.load(f)
+                
+            # Handle both cases: array of users or single object
+            if isinstance(json_data, list):
+                users_data = json_data
+            elif isinstance(json_data, dict):
+                # If it's a single record, put it in a list
+                users_data = [json_data]
+            else:
+                raise ValueError("Unsupported JSON format")
+                
+        except json.JSONDecodeError:
+            # If it fails, the file might be a non-standard JSON (objects separated by commas)
+            logging.info("Detected possible non-standard JSON format, attempting correction...")
+            with open(DATA_PATH, 'r') as f:
+                text = f.read().strip()
+                
+            # Add square brackets to create a valid JSON array
+            if not text.startswith('['):
+                text = '[' + text
+            if not text.endswith(']'):
+                text = text + ']'
+                
+            # Try loading the corrected text
+            users_data = json.loads(text)
+            logging.info("JSON format successfully corrected")
+            
+        logging.info(f"Loaded {len(users_data)} records")
+    except FileNotFoundError:
+        logging.error(f"File {DATA_PATH} not found!")
+        return
+    except Exception as e:
+        logging.error(f"Unable to load file: {str(e)}")
+        return
+
+    # Initialize pipeline
+    logging.info("Initializing pipeline...")
+    pipeline = UserEmbeddingPipeline(
+        output_dim=OUTPUT_DIM,
+        max_sequence_length=MAX_SEQ_LENGTH
+    )
+    
+    # Build vocabularies
+    logging.info("Building vocabularies...")
+    try:
+        pipeline.build_vocabularies(users_data)
+        vocab_sizes = {field: len(vocab) for field, vocab in pipeline.vocab_maps.items()}
+        logging.info(f"Vocabulary sizes: {vocab_sizes}")
+    except Exception as e:
+        logging.error(f"Error building vocabularies: {str(e)}")
+        return
+
+    # Initialize model
+    logging.info("Initializing model...")
+    try:
+        pipeline.initialize_model()
+        logging.info("Model initialized successfully")
+    except Exception as e:
+        logging.error(f"Error initializing model: {str(e)}")
+        return
+
+    # Training
+    logging.info("Starting training...")
+    try:
+        # Create directory for checkpoints if it doesn't exist
+        model_dir = output_dir / "model_checkpoints"
+        model_dir.mkdir(exist_ok=True, parents=True)
+        
+        model = train_user_embeddings(
+            pipeline.model,
+            users_data,
+            pipeline,  # Pass the pipeline to training
+            num_epochs=NUM_EPOCHS,
+            batch_size=BATCH_SIZE,
+            lr=LEARNING_RATE,
+            save_dir=model_dir,  # Add incremental saving
+            save_interval=SAVE_INTERVAL,  # Save every N epochs
+            num_triplets=NUM_TRIPLETS
+        )
+        logging.info("Training completed")
+        pipeline.model = model
+        
+        # Save only the model file
+        logging.info("Saving model...")
+        
+        # Create output directory
+        output_dir.mkdir(exist_ok=True)
+        
+        # Save path
+        model_path = output_dir / 'model.pth'
+        
+        # Prepare dictionary with model state and metadata
+        checkpoint = {
+            'model_state_dict': pipeline.model.state_dict(),
+            'vocab_maps': pipeline.vocab_maps,
+            'embedding_dims': pipeline.embedding_dims,
+            'output_dim': pipeline.output_dim,
+            'max_sequence_length': pipeline.max_sequence_length
+        }
+        
+        # Save model
+        torch.save(checkpoint, model_path)
+        
+        logging.info(f"Model saved to: {model_path}")
+        
+        # Also save a configuration file for reference
+        config_info = {
+            'model_type': 'UserEmbeddingModel',
+            'vocab_sizes': {field: len(vocab) for field, vocab in pipeline.vocab_maps.items()},
+            'embedding_dims': pipeline.embedding_dims,
+            'output_dim': pipeline.output_dim,
+            'max_sequence_length': pipeline.max_sequence_length,
+            'padded_fields': list(pipeline.model.padded_fields),
+            'fields': pipeline.fields
+        }
+        
+        config_path = output_dir / 'model_config.json'
+        with open(config_path, 'w') as f:
+            json.dump(config_info, f, indent=2)
+            
+        logging.info(f"Model configuration saved to: {config_path}")
+        
+        # Only save in HuggingFace format if requested
+        save_hf = os.environ.get("SAVE_HF_FORMAT", "false").lower() == "true"
+        if save_hf:
+            logging.info("Saving in HuggingFace format...")
+            # Save model in HuggingFace format
+            hf_dir = output_dir / 'huggingface'
+            hf_dir.mkdir(exist_ok=True)
+            
+            # Save model in HF format
+            torch.save(pipeline.model.state_dict(), hf_dir / 'pytorch_model.bin')
+            
+            # Save config
+            with open(hf_dir / 'config.json', 'w') as f:
+                json.dump(config_info, f, indent=2)
+                
+            logging.info(f"Model saved in HuggingFace format to: {hf_dir}")
+        
+        # Push to HuggingFace if environment variable is set
+        hf_repo_id = os.environ.get("HF_REPO_ID")
+        hf_token = os.environ.get("HF_TOKEN")
+        
+        if save_hf and hf_repo_id and hf_token:
+            try:
+                from huggingface_hub import HfApi
+                
+                logging.info(f"Pushing model to HuggingFace: {hf_repo_id}")
+                api = HfApi()
+                
+                # Push the model directory
+                api.create_repo(
+                    repo_id=hf_repo_id,
+                    token=hf_token,
+                    exist_ok=True,
+                    private=True
+                )
+                
+                # Upload the files
+                for file_path in (output_dir / "huggingface").glob("**/*"):
+                    if file_path.is_file():
+                        api.upload_file(
+                            path_or_fileobj=str(file_path),
+                            path_in_repo=file_path.relative_to(output_dir / "huggingface"),
+                            repo_id=hf_repo_id,
+                            token=hf_token
+                        )
+                
+                logging.info(f"Model successfully pushed to HuggingFace: {hf_repo_id}")
+            except Exception as e:
+                logging.error(f"Error pushing to HuggingFace: {str(e)}")
+    
+    except Exception as e:
+        logging.error(f"Error during training or saving: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return
+
+    logging.info("Process completed successfully!")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,12 @@
+torch==2.1.0
+torchvision==0.16.0
+numpy==1.24.3
+pandas==2.0.3
+tqdm==4.66.1
+scikit-learn==1.3.0
+matplotlib==3.7.2
+transformers==4.34.0
+datasets==2.14.5
+huggingface_hub==0.17.3
+tensorboard==2.14.1
+psutil==5.9.5

run.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+set -e
+
+echo "Starting user embedding model pipeline..."
+
+# Check if model already exists
+if [ -f "embeddings_output/model.pth" ]; then
+  echo "Model already trained."
+else
+  # Check if input data exists
+  if [ ! -f "$DATA_PATH" ] && [ ! -f "users.json" ]; then
+    echo "Error: No data file found. Please mount a volume with users.json or set DATA_PATH."
+    exit 1
+  fi
+
+  # Run the model generation script
+  
+  python3 generate_model_gpu.py
+  echo "Process completed. Check embeddings_output directory for results."
+fi
+
+# Get the hostname or public URL if available (for Hugging Face Spaces)
+if [ -n "$SPACE_ID" ]; then
+  # If running in Hugging Face Spaces
+  BASE_URL="https://${SPACE_ID}.hf.space"
+else
+  # If running locally or in generic Docker
+  BASE_URL="http://localhost:7860"
+fi
+
+echo "=========================================================="
+echo "Model is available for download at the following URLs:"
+echo "${BASE_URL}/file=embeddings_output/model.pth"
+echo "${BASE_URL}/file=embeddings_output/model_config.json"
+echo "=========================================================="
+
+# Start the web server to serve files
+echo "Starting web server on port 7860..."
+python3 -m http.server 7860

users_200k.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:066da954febcc008311c23735766b751be086b4e61394b797692916f488ccbe4
-size 297966576
+oid sha256:5d91ddb646dfcc955f89edcf756041415a6ba9f3d34cd217e8dc7a6f35cc9161
+size 297966779