be52af70
 tangwang
first commit
|
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
|
"""
Data transformer for converting source data to ES documents.
Handles field mapping, type conversion, and embedding generation.
"""
import pandas as pd
import numpy as np
from typing import Dict, Any, List, Optional
from config import CustomerConfig, FieldConfig, FieldType
from embeddings import BgeEncoder, CLIPImageEncoder
from utils.cache import EmbeddingCache
class DataTransformer:
"""Transform source data into ES-ready documents."""
def __init__(
self,
config: CustomerConfig,
text_encoder: Optional[BgeEncoder] = None,
image_encoder: Optional[CLIPImageEncoder] = None,
use_cache: bool = True
):
"""
Initialize data transformer.
Args:
config: Customer configuration
text_encoder: Text embedding encoder (lazy loaded if not provided)
image_encoder: Image embedding encoder (lazy loaded if not provided)
use_cache: Whether to use embedding cache
"""
self.config = config
self._text_encoder = text_encoder
self._image_encoder = image_encoder
self.use_cache = use_cache
if use_cache:
self.text_cache = EmbeddingCache(".cache/text_embeddings")
self.image_cache = EmbeddingCache(".cache/image_embeddings")
else:
self.text_cache = None
self.image_cache = None
@property
def text_encoder(self) -> BgeEncoder:
"""Lazy load text encoder."""
if self._text_encoder is None:
print("[DataTransformer] Initializing text encoder...")
self._text_encoder = BgeEncoder()
return self._text_encoder
@property
def image_encoder(self) -> CLIPImageEncoder:
"""Lazy load image encoder."""
if self._image_encoder is None:
print("[DataTransformer] Initializing image encoder...")
self._image_encoder = CLIPImageEncoder()
return self._image_encoder
def transform_batch(
self,
df: pd.DataFrame,
batch_size: int = 32
) -> List[Dict[str, Any]]:
"""
Transform a batch of source data into ES documents.
Args:
df: DataFrame with source data
batch_size: Batch size for embedding generation
Returns:
List of ES documents
"""
documents = []
# First pass: generate all embeddings in batch
embedding_data = self._generate_embeddings_batch(df, batch_size)
# Second pass: build documents
for idx, row in df.iterrows():
doc = self._transform_row(row, embedding_data.get(idx, {}))
if doc:
documents.append(doc)
return documents
def _generate_embeddings_batch(
self,
df: pd.DataFrame,
batch_size: int
) -> Dict[int, Dict[str, Any]]:
"""
Generate all embeddings in batch for efficiency.
Args:
df: Source dataframe
batch_size: Batch size
Returns:
Dictionary mapping row index to embedding data
"""
result = {}
# Collect all text embedding fields
text_embedding_fields = [
field for field in self.config.fields
if field.field_type == FieldType.TEXT_EMBEDDING
]
# Collect all image embedding fields
image_embedding_fields = [
field for field in self.config.fields
if field.field_type == FieldType.IMAGE_EMBEDDING
]
# Process text embeddings
for field in text_embedding_fields:
source_col = field.source_column
if source_col not in df.columns:
continue
print(f"[DataTransformer] Generating text embeddings for field: {field.name}")
# Get texts and check cache
texts_to_encode = []
text_indices = []
for idx, row in df.iterrows():
text = row[source_col]
if pd.isna(text) or text == '':
continue
text_str = str(text)
# Check cache
if self.use_cache and self.text_cache.exists(text_str):
cached_emb = self.text_cache.get(text_str)
if idx not in result:
result[idx] = {}
result[idx][field.name] = cached_emb
else:
texts_to_encode.append(text_str)
text_indices.append(idx)
# Encode batch
if texts_to_encode:
embeddings = self.text_encoder.encode_batch(
texts_to_encode,
batch_size=batch_size
)
# Store results
for i, (idx, emb) in enumerate(zip(text_indices, embeddings)):
if idx not in result:
result[idx] = {}
result[idx][field.name] = emb
# Cache
if self.use_cache:
self.text_cache.set(texts_to_encode[i], emb)
# Process image embeddings
for field in image_embedding_fields:
source_col = field.source_column
if source_col not in df.columns:
continue
print(f"[DataTransformer] Generating image embeddings for field: {field.name}")
# Get URLs and check cache
urls_to_encode = []
url_indices = []
for idx, row in df.iterrows():
url = row[source_col]
if pd.isna(url) or url == '':
continue
url_str = str(url)
# Check cache
if self.use_cache and self.image_cache.exists(url_str):
cached_emb = self.image_cache.get(url_str)
if idx not in result:
result[idx] = {}
result[idx][field.name] = cached_emb
else:
urls_to_encode.append(url_str)
url_indices.append(idx)
# Encode batch (with smaller batch size for images)
if urls_to_encode:
embeddings = self.image_encoder.encode_batch(
urls_to_encode,
batch_size=min(8, batch_size)
)
# Store results
for i, (idx, emb) in enumerate(zip(url_indices, embeddings)):
if emb is not None:
if idx not in result:
result[idx] = {}
result[idx][field.name] = emb
# Cache
if self.use_cache:
self.image_cache.set(urls_to_encode[i], emb)
return result
def _transform_row(
self,
row: pd.Series,
embedding_data: Dict[str, Any]
) -> Optional[Dict[str, Any]]:
"""
Transform a single row into an ES document.
Args:
row: Source data row
embedding_data: Pre-computed embeddings for this row
Returns:
ES document or None if transformation fails
"""
doc = {}
for field in self.config.fields:
field_name = field.name
source_col = field.source_column
# Handle embedding fields
if field.field_type in [FieldType.TEXT_EMBEDDING, FieldType.IMAGE_EMBEDDING]:
if field_name in embedding_data:
emb = embedding_data[field_name]
if isinstance(emb, np.ndarray):
doc[field_name] = emb.tolist()
continue
# Handle regular fields
if source_col not in row:
if field.required:
print(f"Warning: Required field '{field_name}' missing in row")
return None
continue
value = row[source_col]
# Skip null values for non-required fields
if pd.isna(value):
if field.required:
print(f"Warning: Required field '{field_name}' is null")
return None
continue
# Type conversion
converted_value = self._convert_value(value, field)
if converted_value is not None:
doc[field_name] = converted_value
return doc
def _convert_value(self, value: Any, field: FieldConfig) -> Any:
"""Convert value to appropriate type for ES."""
if pd.isna(value):
return None
field_type = field.field_type
if field_type == FieldType.TEXT:
return str(value)
elif field_type == FieldType.KEYWORD:
return str(value)
elif field_type in [FieldType.INT, FieldType.LONG]:
try:
return int(value)
except (ValueError, TypeError):
return None
elif field_type in [FieldType.FLOAT, FieldType.DOUBLE]:
try:
return float(value)
except (ValueError, TypeError):
return None
elif field_type == FieldType.BOOLEAN:
if isinstance(value, bool):
return value
if isinstance(value, (int, float)):
return bool(value)
if isinstance(value, str):
return value.lower() in ['true', '1', 'yes', 'y']
return None
elif field_type == FieldType.DATE:
# Pandas datetime handling
if isinstance(value, pd.Timestamp):
return value.isoformat()
|
|
bb3c5ef8
tangwang
灌入数据流程跑通
|
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
|
elif isinstance(value, str):
# Try to parse string datetime and convert to ISO format
try:
import datetime
# Handle common datetime formats
formats = [
'%Y-%m-%d %H:%M:%S', # 2020-07-07 16:44:09
'%Y-%m-%d %H:%M:%S.%f', # 2020-07-07 16:44:09.123
'%Y-%m-%dT%H:%M:%S', # 2020-07-07T16:44:09
'%Y-%m-%d', # 2020-07-07
]
for fmt in formats:
try:
dt = datetime.datetime.strptime(value.strip(), fmt)
return dt.isoformat()
except ValueError:
continue
# If no format matches, return original string
return value
except Exception:
return value
return value
|