FEAT: added server coordination and model aggregation logic

src/main.py

@@ -0,0 +1,62 @@
+import argparse
+import yaml
+import logging
+import logging.config
+from pathlib import Path
+from src.server.coordinator import FederatedCoordinator
+from src.client.model import FederatedClient
+
+def setup_logging(config):
+    """Setup logging configuration."""
+    # Create logs directory if it doesn't exist
+    Path("logs").mkdir(exist_ok=True)
+    
+    log_level = (config.get('monitoring', {}).get('log_level') 
+                or config.get('server', {}).get('monitoring', {}).get('log_level') 
+                or config.get('client', {}).get('monitoring', {}).get('log_level')
+                or 'INFO')
+    
+    # Configure logging with UTF-8 encoding
+    logging.basicConfig(
+        level=log_level,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.StreamHandler(),
+            logging.FileHandler('logs/federated_learning.log', mode='a', encoding='utf-8')
+        ]
+    )
+
+    # Reduce TensorFlow logging noise
+    logging.getLogger('tensorflow').setLevel(logging.WARNING)
+    
+    # Create a divider in the log file
+    logger = logging.getLogger(__name__)
+    logger.info("\n" + "="*50)
+    logger.info("New Training Session Started")
+    logger.info("="*50 + "\n")
+
+def load_config(config_path: str) -> dict:
+    with open(config_path, 'r') as f:
+        return yaml.safe_load(f)
+
+def main():
+    parser = argparse.ArgumentParser(description='Federated Learning Demo')
+    parser.add_argument('--mode', choices=['server', 'client'], required=True)
+    parser.add_argument('--config', type=str, required=True)
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    setup_logging(config)
+    logger = logging.getLogger(__name__)
+
+    if args.mode == 'server':
+        coordinator = FederatedCoordinator(config)
+        logger.info("Starting server...")
+        coordinator.start()
+    else:
+        client = FederatedClient(1, config)
+        logger.info(f"Starting client with ID: {client.client_id}")
+        client.start()
+
+if __name__ == "__main__":
+    main()

src/models/gan.py

@@ -0,0 +1,122 @@
+"""GAN implementation for financial data generation."""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from typing import Dict, Tuple
+
+class Generator(nn.Module):
+    def __init__(self, latent_dim: int, feature_dim: int, hidden_dims: List[int]):
+        super().__init__()
+        
+        layers = []
+        prev_dim = latent_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, hidden_dim),
+                nn.BatchNorm1d(hidden_dim),
+                nn.LeakyReLU(0.2),
+                nn.Dropout(0.3)
+            ])
+            prev_dim = hidden_dim
+            
+        layers.append(nn.Linear(prev_dim, feature_dim))
+        layers.append(nn.Tanh())
+        
+        self.model = nn.Sequential(*layers)
+        
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        return self.model(z)
+
+class Discriminator(nn.Module):
+    def __init__(self, feature_dim: int, hidden_dims: List[int]):
+        super().__init__()
+        
+        layers = []
+        prev_dim = feature_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, hidden_dim),
+                nn.LeakyReLU(0.2),
+                nn.Dropout(0.3)
+            ])
+            prev_dim = hidden_dim
+            
+        layers.append(nn.Linear(prev_dim, 1))
+        layers.append(nn.Sigmoid())
+        
+        self.model = nn.Sequential(*layers)
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+class FinancialGAN:
+    def __init__(self, config: Dict):
+        """Initialize the GAN."""
+        self.latent_dim = config['model']['latent_dim']
+        self.feature_dim = config['model']['feature_dim']
+        self.hidden_dims = config['model']['hidden_dims']
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        self.generator = Generator(
+            self.latent_dim, 
+            self.feature_dim, 
+            self.hidden_dims
+        ).to(self.device)
+        
+        self.discriminator = Discriminator(
+            self.feature_dim, 
+            self.hidden_dims[::-1]
+        ).to(self.device)
+        
+        self.g_optimizer = optim.Adam(
+            self.generator.parameters(), 
+            lr=config['model']['learning_rate']
+        )
+        self.d_optimizer = optim.Adam(
+            self.discriminator.parameters(), 
+            lr=config['model']['learning_rate']
+        )
+        
+        self.criterion = nn.BCELoss()
+        
+    def train_step(self, real_data: torch.Tensor) -> Tuple[float, float]:
+        """Perform one training step."""
+        batch_size = real_data.size(0)
+        real_label = torch.ones(batch_size, 1).to(self.device)
+        fake_label = torch.zeros(batch_size, 1).to(self.device)
+        
+        # Train Discriminator
+        self.d_optimizer.zero_grad()
+        d_real_output = self.discriminator(real_data)
+        d_real_loss = self.criterion(d_real_output, real_label)
+        
+        z = torch.randn(batch_size, self.latent_dim).to(self.device)
+        fake_data = self.generator(z)
+        d_fake_output = self.discriminator(fake_data.detach())
+        d_fake_loss = self.criterion(d_fake_output, fake_label)
+        
+        d_loss = d_real_loss + d_fake_loss
+        d_loss.backward()
+        self.d_optimizer.step()
+        
+        # Train Generator
+        self.g_optimizer.zero_grad()
+        g_output = self.discriminator(fake_data)
+        g_loss = self.criterion(g_output, real_label)
+        g_loss.backward()
+        self.g_optimizer.step()
+        
+        return g_loss.item(), d_loss.item()
+        
+    def generate_samples(self, num_samples: int) -> torch.Tensor:
+        """Generate synthetic financial data."""
+        self.generator.eval()
+        with torch.no_grad():
+            z = torch.randn(num_samples, self.latent_dim).to(self.device)
+            samples = self.generator(z)
+        self.generator.train()
+        return samples
+

src/models/vae.py

@@ -0,0 +1,52 @@
+"""vae.py module."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List
+
+class VAE(nn.Module):
+    def __init__(self, input_dim: int, latent_dim: int, hidden_dims: List[int]):
+        super(VAE, self).__init__()
+        
+        # Encoder
+        modules = []
+        in_features = input_dim
+        for h_dim in hidden_dims:
+            modules.append(nn.Linear(in_features, h_dim))
+            modules.append(nn.ReLU())
+            in_features = h_dim
+        self.encoder = nn.Sequential(*modules)
+        
+        # Latent space
+        self.fc_mu = nn.Linear(hidden_dims[-1], latent_dim)
+        self.fc_var = nn.Linear(hidden_dims[-1], latent_dim)
+        
+        # Decoder
+        modules = []
+        hidden_dims.reverse()
+        in_features = latent_dim
+        for h_dim in hidden_dims:
+            modules.append(nn.Linear(in_features, h_dim))
+            modules.append(nn.ReLU())
+            in_features = h_dim
+        modules.append(nn.Linear(hidden_dims[-1], input_dim))
+        self.decoder = nn.Sequential(*modules)
+        
+    def encode(self, x):
+        h = self.encoder(x)
+        return self.fc_mu(h), self.fc_var(h)
+        
+    def decode(self, z):
+        return self.decoder(z)
+        
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+        
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        return self.decode(z), mu, log_var
+

src/rag/generator.py

@@ -0,0 +1,62 @@
+"""Generator component for the RAG system."""
+
+from typing import List, Dict
+import torch
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM,
+    LogitsProcessor,
+    LogitsProcessorList
+)
+
+class FinancialContextProcessor(LogitsProcessor):
+    """Custom logits processor for financial context."""
+    def __init__(self, financial_constraints: Dict):
+        self.constraints = financial_constraints
+        
+    def __call__(self, input_ids: torch.LongTensor, 
+                 scores: torch.FloatTensor) -> torch.FloatTensor:
+        # Apply financial domain constraints
+        # This is a placeholder for actual constraints
+        return scores
+
+class RAGGenerator:
+    def __init__(self, config: Dict):
+        """Initialize the generator."""
+        self.model_name = "gpt2"  # Can be configured based on needs
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
+        self.max_length = 512
+        
+    def prepare_context(self, retrieved_docs: List[Dict]) -> str:
+        """Prepare context from retrieved documents."""
+        context = ""
+        for doc in retrieved_docs:
+            context += f"{doc['document']['text']}\n"
+        return context.strip()
+        
+    def generate(self, query: str, retrieved_docs: List[Dict], 
+                financial_constraints: Dict = None) -> str:
+        """Generate text based on query and retrieved documents."""
+        context = self.prepare_context(retrieved_docs)
+        prompt = f"Context: {context}\nQuery: {query}\nResponse:"
+        
+        # Prepare logits processors
+        processors = LogitsProcessorList()
+        if financial_constraints:
+            processors.append(FinancialContextProcessor(financial_constraints))
+            
+        # Generate response
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        outputs = self.model.generate(
+            inputs.input_ids,
+            max_length=self.max_length,
+            num_return_sequences=1,
+            logits_processor=processors,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )
+        
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+

src/rag/retriever.py

@@ -0,0 +1,84 @@
+"""Retrieval component for the RAG system."""
+
+import faiss
+import numpy as np
+from typing import List, Dict, Tuple
+from elasticsearch import Elasticsearch
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+class FinancialDataRetriever:
+    def __init__(self, config: Dict):
+        """Initialize the retriever with configuration."""
+        self.retriever_type = config['rag']['retriever']
+        self.max_documents = config['rag']['max_documents']
+        self.similarity_threshold = config['rag']['similarity_threshold']
+        
+        # Initialize FAISS index
+        self.dimension = 768  # BERT embedding dimension
+        self.index = faiss.IndexFlatL2(self.dimension)
+        
+        # Initialize transformer model for embeddings
+        self.tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
+        self.model = AutoModel.from_pretrained('bert-base-uncased')
+        
+        # Initialize Elasticsearch if needed
+        if self.retriever_type == "elasticsearch":
+            self.es = Elasticsearch()
+            
+    def encode_text(self, texts: List[str]) -> np.ndarray:
+        """Encode text using BERT."""
+        tokens = self.tokenizer(texts, padding=True, truncation=True, 
+                              return_tensors="pt", max_length=512)
+        with torch.no_grad():
+            outputs = self.model(**tokens)
+            embeddings = outputs.last_hidden_state[:, 0, :].numpy()
+        return embeddings
+        
+    def index_documents(self, documents: List[Dict]):
+        """Index documents for retrieval."""
+        if self.retriever_type == "faiss":
+            texts = [doc['text'] for doc in documents]
+            embeddings = self.encode_text(texts)
+            self.index.add(embeddings)
+            self.documents = documents
+        else:
+            for doc in documents:
+                self.es.index(index="financial_data", document=doc)
+                
+    def retrieve(self, query: str, k: int = None) -> List[Dict]:
+        """Retrieve relevant documents."""
+        k = k or self.max_documents
+        query_embedding = self.encode_text([query])
+        
+        if self.retriever_type == "faiss":
+            distances, indices = self.index.search(query_embedding, k)
+            results = [
+                {
+                    'document': self.documents[idx],
+                    'score': float(1 / (1 + dist))
+                }
+                for dist, idx in zip(distances[0], indices[0])
+                if 1 / (1 + dist) >= self.similarity_threshold
+            ]
+        else:
+            response = self.es.search(
+                index="financial_data",
+                query={
+                    "match": {
+                        "text": query
+                    }
+                },
+                size=k
+            )
+            results = [
+                {
+                    'document': hit['_source'],
+                    'score': hit['_score']
+                }
+                for hit in response['hits']['hits']
+                if hit['_score'] >= self.similarity_threshold
+            ]
+            
+        return results
+

src/server/aggregator.py

@@ -0,0 +1,43 @@
+"""aggregator.py module."""
+
+import tensorflow as tf
+from typing import List, Dict
+import numpy as np
+from collections import defaultdict
+
+class FederatedAggregator:
+    def __init__(self, config: Dict):
+        """Initialize the federated aggregator."""
+        self.weighted = config['aggregation']['weighted']
+        
+    def compute_metrics(self, client_metrics: List[Dict]) -> Dict:
+        """Compute aggregated metrics from client updates."""
+        if not client_metrics:
+            return {}
+            
+        aggregated_metrics = defaultdict(float)
+        total_samples = sum(metrics['num_samples'] for metrics in client_metrics)
+        
+        for metrics in client_metrics:
+            weight = metrics['num_samples'] / total_samples if self.weighted else 1.0
+            
+            for metric_name, value in metrics['metrics'].items():
+                aggregated_metrics[metric_name] += value * weight
+                
+        return dict(aggregated_metrics)
+    
+    def check_convergence(self, 
+                         old_weights: List, 
+                         new_weights: List, 
+                         threshold: float = 1e-5) -> bool:
+        """Check if the model has converged."""
+        if old_weights is None or new_weights is None:
+            return False
+            
+        weight_differences = [
+            np.mean(np.abs(old - new))
+            for old, new in zip(old_weights, new_weights)
+        ]
+        
+        return all(diff < threshold for diff in weight_differences)
+

src/server/coordinator.py

@@ -0,0 +1,77 @@
+"""coordinator.py module."""
+
+import tensorflow as tf
+from typing import List, Dict
+import numpy as np
+from collections import defaultdict
+import logging
+import time
+
+class FederatedCoordinator:
+    def __init__(self, config: Dict):
+        """Initialize the federated learning coordinator."""
+        self.config = config
+        self.clients = {}
+        self.current_round = 0
+        self.min_clients = config.get('server', {}).get('federated', {}).get('min_clients', 2)
+        self.rounds = config.get('server', {}).get('federated', {}).get('rounds', 10)
+        
+    def register_client(self, client_id: int, client_size: int):
+        """Register a new client."""
+        self.clients[client_id] = {
+            'size': client_size,
+            'weights': None,
+            'metrics': defaultdict(list)
+        }
+    
+    def aggregate_weights(self, client_updates: List[Dict]) -> List:
+        """Aggregate weights using FedAvg algorithm."""
+        total_size = sum(self.clients[update['client_id']]['size'] 
+                        for update in client_updates)
+        
+        aggregated_weights = [
+            np.zeros_like(w) for w in client_updates[0]['weights']
+        ]
+        
+        for update in client_updates:
+            client_size = self.clients[update['client_id']]['size']
+            weight = client_size / total_size
+            
+            for i, layer_weights in enumerate(update['weights']):
+                aggregated_weights[i] += layer_weights * weight
+                
+        return aggregated_weights
+    
+    def start(self):
+        """Start the federated learning process."""
+        logger = logging.getLogger(__name__)
+        
+        # Print server startup information
+        logger.info("\n" + "=" * 60)
+        logger.info(f"{'Federated Learning Server Starting':^60}")
+        logger.info("=" * 60)
+        
+        # Print configuration details
+        logger.info("\nServer Configuration:")
+        logger.info("-" * 30)
+        logger.info(f"Minimum clients required: {self.min_clients}")
+        logger.info(f"Total rounds planned: {self.rounds}")
+        logger.info(f"Current active clients: {len(self.clients)}")
+        logger.info("-" * 30 + "\n")
+        
+        while self.current_round < self.rounds:
+            round_num = self.current_round + 1
+            logger.info(f"\nRound {round_num}/{self.rounds}")
+            logger.info("-" * 30)
+            
+            if len(self.clients) < self.min_clients:
+                logger.warning(
+                    f"Waiting for clients... "
+                    f"(active: {len(self.clients)}/{self.min_clients})"
+                )
+                time.sleep(5)
+                continue
+            
+            logger.info(f"Active clients: {list(self.clients.keys())}")
+            logger.info(f"Starting training round {round_num}")
+            self.current_round += 1

src/utils/metrics.py

@@ -0,0 +1,119 @@
+"""metrics.py module."""
+
+from typing import Dict, List
+import numpy as np
+from scipy.stats import wasserstein_distance, ks_2samp
+from sklearn.metrics import mutual_info_score, silhouette_score
+from sklearn.neighbors import NearestNeighbors
+
+class MetricsCalculator:
+    @staticmethod
+    def calculate_distribution_similarity(real_data: np.ndarray, 
+                                       synthetic_data: np.ndarray) -> Dict[str, float]:
+        """Calculate statistical similarity metrics between real and synthetic data."""
+        metrics = {}
+        
+        # Wasserstein distance
+        metrics['wasserstein'] = wasserstein_distance(
+            real_data.flatten(), 
+            synthetic_data.flatten()
+        )
+        
+        # KL divergence approximation
+        metrics['mutual_info'] = mutual_info_score(
+            real_data.flatten(),
+            synthetic_data.flatten()
+        )
+        
+        # Kolmogorov-Smirnov test
+        ks_statistic, p_value = ks_2samp(real_data.flatten(), synthetic_data.flatten())
+        metrics['ks_statistic'] = ks_statistic
+        metrics['ks_p_value'] = p_value
+        
+        # Basic statistical measures
+        metrics['mean_diff'] = abs(np.mean(real_data) - np.mean(synthetic_data))
+        metrics['std_diff'] = abs(np.std(real_data) - np.std(synthetic_data))
+        metrics['percentile_diff'] = np.mean([
+            abs(np.percentile(real_data, p) - np.percentile(synthetic_data, p))
+            for p in [25, 50, 75]
+        ])
+        
+        return metrics
+    
+    @staticmethod
+    def evaluate_privacy_metrics(model, test_data: np.ndarray, 
+                               synthetic_data: np.ndarray) -> Dict[str, float]:
+        """Evaluate privacy-related metrics."""
+        metrics = {}
+        
+        # Membership inference risk
+        metrics['membership_inference_risk'] = MetricsCalculator._calculate_membership_inference_risk(
+            test_data, synthetic_data
+        )
+        
+        # Attribute inference risk
+        metrics['attribute_inference_risk'] = MetricsCalculator._calculate_attribute_inference_risk(
+            test_data, synthetic_data
+        )
+        
+        # k-anonymity approximation
+        metrics['k_anonymity_score'] = MetricsCalculator._calculate_k_anonymity(synthetic_data)
+        
+        # Uniqueness score
+        metrics['uniqueness_score'] = MetricsCalculator._calculate_uniqueness(synthetic_data)
+        
+        return metrics
+    
+    @staticmethod
+    def _calculate_membership_inference_risk(test_data: np.ndarray, 
+                                          synthetic_data: np.ndarray) -> float:
+        """Calculate membership inference risk using nearest neighbor distance ratio."""
+        k = 3  # number of neighbors to consider
+        nn = NearestNeighbors(n_neighbors=k)
+        nn.fit(synthetic_data)
+        
+        distances, _ = nn.kneighbors(test_data)
+        avg_min_distances = distances.mean(axis=1)
+        
+        # Normalize to [0,1] where higher values indicate higher privacy
+        risk_score = 1.0 - (1.0 / (1.0 + np.mean(avg_min_distances)))
+        return risk_score
+    
+    @staticmethod
+    def _calculate_attribute_inference_risk(test_data: np.ndarray, 
+                                         synthetic_data: np.ndarray) -> float:
+        """Calculate attribute inference risk using correlation analysis."""
+        real_corr = np.corrcoef(test_data.T)
+        synth_corr = np.corrcoef(synthetic_data.T)
+        
+        # Compare correlation matrices
+        correlation_diff = np.abs(real_corr - synth_corr).mean()
+        
+        # Convert to risk score (0 to 1, where lower is better)
+        risk_score = 1.0 - np.exp(-correlation_diff)
+        return risk_score
+    
+    @staticmethod
+    def _calculate_k_anonymity(data: np.ndarray, k: int = 5) -> float:
+        """Calculate approximate k-anonymity score."""
+        nn = NearestNeighbors(n_neighbors=k)
+        nn.fit(data)
+        
+        distances, _ = nn.kneighbors(data)
+        k_anonymity_scores = distances[:, -1]  # Distance to k-th neighbor
+        
+        # Convert to score (0 to 1, where higher is better)
+        return float(np.mean(k_anonymity_scores > 0.1))
+    
+    @staticmethod
+    def _calculate_uniqueness(data: np.ndarray) -> float:
+        """Calculate uniqueness score of the dataset."""
+        nn = NearestNeighbors(n_neighbors=2)  # 2 because first neighbor is self
+        nn.fit(data)
+        
+        distances, _ = nn.kneighbors(data)
+        uniqueness_scores = distances[:, 1]  # Distance to nearest non-self neighbor
+        
+        # Convert to score (0 to 1, where higher means more unique records)
+        return float(np.mean(uniqueness_scores > np.median(uniqueness_scores)))
+

src/utils/privacy.py

@@ -0,0 +1,37 @@
+"""privacy.py module."""
+
+import tensorflow_privacy as tfp
+from typing import Dict, Any
+import numpy as np
+
+class PrivacyManager:
+    def __init__(self, config: Dict[str, Any]):
+        self.epsilon = config['privacy']['epsilon']
+        self.delta = config['privacy']['delta']
+        self.noise_multiplier = config['privacy']['noise_multiplier']
+        
+    def add_noise_to_gradients(self, gradients: np.ndarray) -> np.ndarray:
+        """Add Gaussian noise to gradients for differential privacy."""
+        noise = np.random.normal(0, self.noise_multiplier, gradients.shape)
+        return gradients + noise
+        
+    def verify_privacy_budget(self, num_iterations: int) -> bool:
+        """Check if training stays within privacy budget."""
+        eps = self.compute_epsilon(num_iterations)
+        return eps <= self.epsilon
+        
+    def compute_epsilon(self, num_iterations: int) -> float:
+        """Compute the current epsilon value."""
+        q = 1.0  # sampling ratio
+        steps = num_iterations
+        orders = ([1.25, 1.5, 1.75, 2., 2.25, 2.5, 3., 3.5, 4., 4.5] +
+                 list(range(5, 64)) + [128, 256, 512])
+        
+        return tfp.compute_dp_sgd_privacy(
+            n=1000,  # number of training points
+            batch_size=32,
+            noise_multiplier=self.noise_multiplier,
+            epochs=steps,
+            delta=self.delta
+        )[0]
+

tests/test_client.py

@@ -0,0 +1,49 @@
+"""test_client.py module."""
+
+import pytest
+import tensorflow as tf
+import yaml
+from src.client.data_handler import FinancialDataHandler
+from src.client.model import FederatedClient
+
+@pytest.fixture
+def config():
+    """Load test configuration."""
+    with open('config/client_config.yaml', 'r') as f:
+        return yaml.safe_load(f)['client']
+
+def test_data_handler(config):
+    """Test data handler functionality."""
+    handler = FinancialDataHandler(config)
+    
+    # Test data simulation
+    data = handler.simulate_financial_data(num_samples=100)
+    assert len(data) == 100
+    assert all(col in data.columns for col in [
+        'transaction_amount',
+        'account_balance',
+        'transaction_frequency',
+        'credit_score',
+        'days_since_last_transaction'
+    ])
+    
+    # Test preprocessing
+    dataset, scaler = handler.get_client_data()
+    assert isinstance(dataset, tf.data.Dataset)
+    
+def test_federated_client(config):
+    """Test federated client functionality."""
+    client = FederatedClient(config)
+    
+    # Test model building
+    assert isinstance(client.model, tf.keras.Model)
+    
+    # Test local training
+    handler = FinancialDataHandler(config)
+    dataset, _ = handler.get_client_data()
+    
+    training_result = client.train_local_model(dataset, epochs=1)
+    assert 'client_id' in training_result
+    assert 'weights' in training_result
+    assert 'metrics' in training_result
+