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
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335"""Training data converter for node type prediction."""
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass, field
import json
import pickle
from collections import defaultdict
from src.milestone1 import (
DSLLoader,
GraphAnalyzer,
NodeTypeMapper,
SequenceExtractor
)
@dataclass
class TrainingSequence:
"""Represents a single training sequence."""
sequence_id: str
source_file: str
node_types: List[str]
type_ids: List[int]
length: int
def to_dict(self) -> dict:
"""Convert to dictionary for serialization."""
return {
'sequence_id': self.sequence_id,
'source_file': self.source_file,
'node_types': self.node_types,
'type_ids': self.type_ids,
'length': self.length
}
@dataclass
class TrainingDataset:
"""Collection of training sequences with metadata."""
raw_sequences: List[TrainingSequence] = field(default_factory=list)
type_mapper: NodeTypeMapper = field(default_factory=NodeTypeMapper)
metadata: Dict = field(default_factory=dict)
sequences: List[List[int]] = field(default_factory=list)
labels: List[int] = field(default_factory=list)
@property
def num_sequences(self) -> int:
"""Total number of sequences."""
return len(self.raw_sequences)
@property
def num_types(self) -> int:
"""Number of unique node types."""
return self.type_mapper.num_types
def get_vocab_size(self) -> int:
"""Get vocabulary size for model."""
return self.type_mapper.num_types + 3 # Add special tokens
def add_sequence(self, sequence: TrainingSequence):
"""Add a sequence to the dataset."""
self.raw_sequences.append(sequence)
# Generate training samples from sequence
for i in range(1, len(sequence.type_ids)):
self.sequences.append(sequence.type_ids[:i])
self.labels.append(sequence.type_ids[i])
def get_statistics(self) -> dict:
"""Get dataset statistics."""
if not self.raw_sequences:
return {
'num_sequences': 0,
'num_types': 0,
'avg_length': 0,
'min_length': 0,
'max_length': 0
}
lengths = [seq.length for seq in self.raw_sequences]
return {
'num_sequences': self.num_sequences,
'num_types': self.num_types,
'avg_length': sum(lengths) / len(lengths),
'min_length': min(lengths),
'max_length': max(lengths),
'total_tokens': sum(lengths),
'unique_files': len(set(seq.source_file for seq in self.raw_sequences))
}
def save(self, path: Path):
"""Save dataset to file."""
data = {
'sequences': [seq.to_dict() for seq in self.raw_sequences],
'type_mapping': self.type_mapper.node_type_to_id,
'reverse_mapping': self.type_mapper.id_to_node_type,
'metadata': self.metadata,
'statistics': self.get_statistics()
}
if path.suffix == '.json':
with open(path, 'w') as f:
json.dump(data, f, indent=2)
elif path.suffix == '.pkl':
with open(path, 'wb') as f:
pickle.dump(data, f)
else:
raise ValueError(f"Unsupported file format: {path.suffix}")
def load(self, path: str):
"""Load dataset from file."""
path = Path(path)
if path.suffix == '.json':
with open(path, 'r') as f:
data = json.load(f)
elif path.suffix == '.pkl':
with open(path, 'rb') as f:
data = pickle.load(f)
else:
raise ValueError(f"Unsupported file format: {path.suffix}")
# Restore type mapper
self.type_mapper.node_type_to_id = data['type_mapping']
self.type_mapper.id_to_node_type = {int(k): v for k, v in data['reverse_mapping'].items()}
self.type_mapper.next_id = max(map(int, data['reverse_mapping'].keys())) + 1 if data['reverse_mapping'] else 0
# Restore sequences
self.raw_sequences = []
self.sequences = []
self.labels = []
for seq_data in data['sequences']:
sequence = TrainingSequence(
sequence_id=seq_data['sequence_id'],
source_file=seq_data['source_file'],
node_types=seq_data['node_types'],
type_ids=seq_data['type_ids'],
length=seq_data['length']
)
self.add_sequence(sequence)
self.metadata = data.get('metadata', {})
class TrainingDataConverter:
"""Converts DSL files to training data."""
def __init__(self, type_mapper: Optional[NodeTypeMapper] = None):
"""
Initialize converter.
Args:
type_mapper: Optional shared type mapper
"""
self.type_mapper = type_mapper or NodeTypeMapper()
self.extractor = SequenceExtractor(self.type_mapper)
def convert_file(
self,
file_path: Path,
max_sequence_length: Optional[int] = None
) -> List[TrainingSequence]:
"""
Convert a single DSL file to training sequences.
Args:
file_path: Path to DSL file
max_sequence_length: Optional max length for sequences
Returns:
List of TrainingSequence objects
"""
sequences = []
try:
# Load and analyze DSL
dsl = DSLLoader.load_file(file_path)
analyzer = GraphAnalyzer(dsl)
# Extract raw sequences
raw_sequences = self.extractor.extract_all_sequences(
analyzer,
max_length=max_sequence_length
)
# Convert to training sequences
for i, raw_seq in enumerate(raw_sequences):
training_seq = TrainingSequence(
sequence_id=f"{file_path.stem}_{i}",
source_file=str(file_path),
node_types=raw_seq.node_types,
type_ids=raw_seq.type_ids,
length=len(raw_seq)
)
sequences.append(training_seq)
except Exception as e:
print(f"Error processing {file_path}: {e}")
return sequences
def convert_directory(
self,
directory: Path,
pattern: str = "*.yml",
max_sequence_length: Optional[int] = None,
limit: Optional[int] = None
) -> TrainingDataset:
"""
Convert all DSL files in a directory to training data.
Args:
directory: Directory containing DSL files
pattern: File pattern to match
max_sequence_length: Optional max length for sequences
limit: Optional limit on number of files to process
Returns:
TrainingDataset object
"""
dataset = TrainingDataset(type_mapper=self.type_mapper)
files = list(directory.glob(pattern))
if limit:
files = files[:limit]
for file_path in files:
sequences = self.convert_file(file_path, max_sequence_length)
for seq in sequences:
dataset.add_sequence(seq)
# Add metadata
dataset.metadata = {
'source_directory': str(directory),
'pattern': pattern,
'max_sequence_length': max_sequence_length,
'num_files_processed': len(files)
}
return dataset
def create_training_samples(
self,
dataset: TrainingDataset,
min_context_length: int = 1
) -> List[Tuple[List[int], int]]:
"""
Create training samples from sequences.
For each sequence A-B-C-D, create samples:
- [A] -> B
- [A, B] -> C
- [A, B, C] -> D
Args:
dataset: TrainingDataset object
min_context_length: Minimum context length
Returns:
List of (context, target) tuples
"""
samples = []
for sequence in dataset.raw_sequences:
# Skip sequences that are too short
if sequence.length <= min_context_length:
continue
# Create samples from sequence
for i in range(min_context_length, sequence.length):
context = sequence.type_ids[:i]
target = sequence.type_ids[i]
samples.append((context, target))
return samples
def create_training_batches(
self,
samples: List[Tuple[List[int], int]],
handle_branches: bool = True
) -> Dict[str, List]:
"""
Organize training samples and handle branches.
Args:
samples: List of (context, target) tuples
handle_branches: Whether to handle multiple targets
Returns:
Dictionary with processed training data
"""
if not handle_branches:
# Simple case: no branch handling
return {
'contexts': [s[0] for s in samples],
'targets': [s[1] for s in samples]
}
# Handle branches: group by context
context_targets = defaultdict(list)
for context, target in samples:
context_key = tuple(context)
context_targets[context_key].append(target)
# Process grouped data
processed_contexts = []
processed_targets = []
processed_distributions = []
for context, targets in context_targets.items():
processed_contexts.append(list(context))
# Count target frequencies
target_counts = defaultdict(int)
for target in targets:
target_counts[target] += 1
# Create probability distribution
total = sum(target_counts.values())
distribution = {}
for target, count in target_counts.items():
distribution[target] = count / total
processed_targets.append(list(target_counts.keys()))
processed_distributions.append(distribution)
return {
'contexts': processed_contexts,
'targets': processed_targets,
'distributions': processed_distributions,
'num_branches': sum(1 for t in processed_targets if len(t) > 1)
}