File size: 11,720 Bytes
c8b77b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
"""
Ethics Training Module for MangoMAS Local

This module implements specialized training for ethical reasoning capabilities,
adapted from the AWS backup system for local training.
"""

import json
import logging
import os
import random
from typing import Any, Dict, List

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset

from ..core_framework import SpecializedTrainingModule, TrainingModuleConfig

logger = logging.getLogger(__name__)


class EthicsDataset(Dataset):
    """Dataset for training ethical reasoning capabilities."""

    def __init__(self, data_path: str, tokenizer, max_length: int = 768):
        """
        Initialize the ethics dataset.

        Args:
            data_path: Path to the ethics data file
            tokenizer: Tokenizer for text processing
            max_length: Maximum sequence length
        """
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = self._load_data(data_path)

        logger.info(f"Loaded ethics dataset with {len(self.data)} examples")

    def _load_data(self, data_path: str) -> List[Dict]:
        """Load ethics training data."""
        data = []
        with open(data_path, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    item = json.loads(line.strip())
                    # Validate required fields for ethics data
                    if (
                        "scenario" in item
                        and "ethical_analysis" in item
                        and "recommendation" in item
                    ):
                        data.append(item)
                except json.JSONDecodeError:
                    continue
        return data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]

        # Format the ethics prompt
        prompt = f"Scenario: {item['scenario']}\nEthical Analysis: {item['ethical_analysis']}\nRecommendation: {item['recommendation']}"

        # Tokenize
        encoding = self.tokenizer(
            prompt,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )

        return {
            "input_ids": encoding["input_ids"].squeeze(),
            "attention_mask": encoding["attention_mask"].squeeze(),
            "labels": encoding["input_ids"].squeeze(),
        }


class EthicsTrainingModule(SpecializedTrainingModule):
    """Specialized training module for ethical reasoning capabilities."""

    def __init__(self, config: TrainingModuleConfig, tokenizer):
        """
        Initialize the ethics training module.

        Args:
            config: Module configuration
            tokenizer: Tokenizer for text processing
        """
        super().__init__(config, tokenizer)

        # Initialize ethics-specific components
        self.ethics_loss = nn.CrossEntropyLoss(ignore_index=-100)
        self.metrics = {
            "ethics_loss": 0.0,
            "ethical_consistency": 0.0,
            "principle_alignment": 0.0,
        }

        logger.info("Initialized EthicsTrainingModule")

    def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Prepare a batch of data for ethics training.

        Args:
            batch: The input batch from the dataloader

        Returns:
            Processed batch ready for ethics training
        """
        # Move batch to device
        prepared_batch = {}
        for key, value in batch.items():
            if isinstance(value, torch.Tensor):
                prepared_batch[key] = value.to(self.device)
            else:
                prepared_batch[key] = value

        return prepared_batch

    def compute_loss(
        self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor]
    ) -> torch.Tensor:
        """
        Compute the ethics-specific loss.

        Args:
            student_outputs: Outputs from the student model
            teacher_outputs: Outputs from the teacher model
            batch: The processed input batch

        Returns:
            Loss tensor for ethics training
        """
        try:
            # Extract logits from model outputs
            if hasattr(student_outputs, "logits"):
                student_logits = student_outputs.logits
            else:
                student_logits = student_outputs

            if hasattr(teacher_outputs, "logits"):
                teacher_logits = teacher_outputs.logits
            else:
                teacher_logits = teacher_outputs

            # Get labels from batch
            labels = batch.get("labels", batch.get("input_ids"))

            # Compute cross entropy loss for ethics
            shift_logits = student_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()

            ethics_loss = self.ethics_loss(
                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
            )

            # Add KL divergence loss between student and teacher
            if teacher_logits is not None:
                kl_loss = F.kl_div(
                    F.log_softmax(student_logits, dim=-1),
                    F.softmax(teacher_logits, dim=-1),
                    reduction="batchmean",
                )
                total_loss = ethics_loss + 0.1 * kl_loss
            else:
                total_loss = ethics_loss

            # Update metrics
            self.metrics["ethics_loss"] = ethics_loss.item()

            return total_loss * self.loss_weight

        except Exception as e:
            logger.error(f"Error computing ethics loss: {e}")
            # Return a small loss to avoid training failure
            return torch.tensor(0.01, requires_grad=True)

    def get_metrics(self) -> Dict[str, float]:
        """
        Get metrics specific to ethics training.

        Returns:
            Dictionary of ethics metrics
        """
        return self.metrics.copy()

    def generate_synthetic_ethics_data(
        self, output_path: str, num_samples: int = 1000
    ) -> None:
        """
        Generate synthetic ethics training data.

        Args:
            output_path: Path to save the generated data
            num_samples: Number of samples to generate
        """
        # This is a simplified implementation based on the AWS backup
        # In a full implementation, this would be much more sophisticated

        ethics_templates = [
            {
                "scenario": "A company wants to collect user location data to improve their mapping service.",
                "principles": [
                    "privacy",
                    "consent",
                    "transparency",
                    "data minimization",
                ],
                "considerations": [
                    "Users should be clearly informed about what data is collected",
                    "Data collection should be opt-in rather than opt-out",
                    "Only necessary data should be collected and stored",
                    "Data should be properly secured and anonymized where possible",
                ],
                "ethical_analysis": "The collection of location data raises privacy concerns, but can be ethical if done with transparency, informed consent, and data minimization practices. Users must be clearly informed about what data is collected, how it's used, and given genuine choice in the matter.",
                "recommendation": "Proceed with location data collection only with explicit opt-in consent, clear privacy notices, data minimization practices, and strong security measures.",
                "stance": "neutral",
            },
            {
                "scenario": "An AI company is developing a facial recognition system to be sold to law enforcement agencies without oversight mechanisms.",
                "principles": [
                    "privacy",
                    "justice",
                    "accountability",
                    "potential for discrimination",
                ],
                "considerations": [
                    "Facial recognition has known bias issues across different demographics",
                    "Law enforcement use creates significant civil liberties concerns",
                    "Lack of oversight could lead to misuse and privacy violations",
                    "Potential chilling effect on free speech and assembly",
                ],
                "ethical_analysis": "Deploying facial recognition in law enforcement without oversight mechanisms raises serious ethical concerns. These systems have demonstrated bias across demographic groups, potentially leading to discriminatory outcomes. Without accountability measures, there's significant risk of misuse, privacy violations, and erosion of civil liberties.",
                "recommendation": "Do not deploy facial recognition systems to law enforcement without robust oversight, accuracy testing across demographics, clear usage limitations, and strong accountability mechanisms.",
                "stance": "harmful",
            },
            {
                "scenario": "A medical AI is being developed to help doctors identify potential early signs of cancer in medical images.",
                "principles": [
                    "beneficence",
                    "non-maleficence",
                    "human oversight",
                    "transparency",
                ],
                "considerations": [
                    "Early detection could save many lives",
                    "False positives could cause unnecessary stress and procedures",
                    "False negatives could delay critical treatment",
                    "System should augment rather than replace medical expertise",
                ],
                "ethical_analysis": "A medical AI for cancer detection has significant potential benefits in improving early diagnosis and saving lives. However, it's critical that the system maintains high accuracy to minimize both false positives (causing unnecessary procedures) and false negatives (missing actual cases). The system should be designed to augment rather than replace medical professionals, with humans making final decisions.",
                "recommendation": "Proceed with development with rigorous clinical validation, transparent reporting of accuracy metrics across diverse populations, clear communication about limitations, and implementation as a decision support tool rather than autonomous system.",
                "stance": "beneficial",
            },
        ]

        # Generate variations
        output_data = []
        for _ in range(num_samples):
            template = random.choice(ethics_templates)

            # Create a minor variation to avoid exact duplicates
            variation = template.copy()

            # Add metadata
            variation["metadata"] = {
                "generated": True,
                "timestamp": (
                    torch.cuda.get_device_name(0)
                    if torch.cuda.is_available()
                    else "CPU"
                ),
            }

            output_data.append(variation)

        # Save to file
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, "w", encoding="utf-8") as f:
            for item in output_data:
                f.write(json.dumps(item) + "\n")

        logger.info(
            f"Generated {len(output_data)} synthetic ethics examples at {output_path}"
        )