eBlessings commited on
Commit
98db640
·
verified ·
1 Parent(s): f7e8c65

Upload T4a.py

Browse files
Files changed (1) hide show
  1. T4a.py +637 -0
T4a.py ADDED
@@ -0,0 +1,637 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ app.py – Quranic Data Training Pipeline Endpoint for T4 Medium
4
+ ----------------------------------------------------------------
5
+ Updated for T4 medium (8 vCores, 30 GB RAM, 16 GB VRAM) with FP16 training,
6
+ checkpoint saving with download link, and enhanced error handling.
7
+ """
8
+
9
+ import json
10
+ import logging
11
+ import os
12
+ import traceback
13
+ import gc
14
+ import time
15
+ import psutil
16
+ import math
17
+ import shutil
18
+ from datetime import datetime
19
+ from typing import Dict, List, Optional
20
+ from dataclasses import dataclass, asdict
21
+
22
+ import torch
23
+ torch.set_num_threads(8)
24
+
25
+ from torch.utils.data import Dataset
26
+ from transformers import (
27
+ AutoTokenizer,
28
+ AutoModelForCausalLM,
29
+ TrainingArguments,
30
+ Trainer,
31
+ DataCollatorForLanguageModeling,
32
+ __version__ as transformers_version
33
+ )
34
+ from threading import Lock
35
+ from accelerate import Accelerator
36
+
37
+ import gradio as gr
38
+ import spaces
39
+
40
+ # Set an environment variable to help mitigate CUDA allocator fragmentation issues.
41
+ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
42
+
43
+ # Updated version requirements
44
+ MIN_TRANSFORMERS_VERSION = "4.45.0"
45
+ if tuple(map(int, transformers_version.split("."))) < tuple(map(int, MIN_TRANSFORMERS_VERSION.split("."))):
46
+ logging.warning(f"Transformers version {transformers_version} detected. Please upgrade to at least {MIN_TRANSFORMERS_VERSION}")
47
+
48
+ logging.basicConfig(
49
+ level=logging.INFO,
50
+ format='%(asctime)s - %(levelname)s - %(message)s',
51
+ handlers=[
52
+ logging.FileHandler('pipeline.log'),
53
+ logging.StreamHandler()
54
+ ]
55
+ )
56
+ logger = logging.getLogger(__name__)
57
+
58
+ def manage_memory(threshold_percent: int = 90, min_available_mb: int = 500, sleep_duration: int = 10):
59
+ vm = psutil.virtual_memory()
60
+ used_percent = vm.percent
61
+ available_mb = vm.available / (1024 * 1024)
62
+ logger.info(f"Memory usage: {used_percent}% used, {available_mb:.2f} MB available")
63
+ if used_percent > threshold_percent or available_mb < min_available_mb:
64
+ logger.warning("High memory usage detected, forcing garbage collection and sleeping...")
65
+ gc.collect()
66
+ time.sleep(sleep_duration)
67
+
68
+ def manage_gpu_resources(sleep_duration: int = 5):
69
+ if torch.cuda.is_available():
70
+ allocated = torch.cuda.memory_allocated() / (1024 * 1024)
71
+ cached = torch.cuda.memory_reserved() / (1024 * 1024)
72
+ logger.info(f"GPU Memory Allocated: {allocated:.2f} MB, Reserved: {cached:.2f} MB")
73
+ torch.cuda.empty_cache()
74
+ time.sleep(sleep_duration)
75
+
76
+ def zip_checkpoint(checkpoint_dir: str) -> str:
77
+ zip_file = checkpoint_dir + ".zip"
78
+ if os.path.exists(zip_file):
79
+ os.remove(zip_file)
80
+ shutil.make_archive(checkpoint_dir, 'zip', checkpoint_dir)
81
+ return os.path.basename(zip_file)
82
+
83
+ @dataclass
84
+ class WordAnalysis:
85
+ arabic: str
86
+ translation: str
87
+ position: str
88
+ morphology: Dict
89
+ features: List[str]
90
+ root: str
91
+ location: str
92
+ metadata: Dict
93
+
94
+ @dataclass
95
+ class VerseData:
96
+ chapter: int
97
+ verse: int
98
+ arabic_text: str
99
+ translation: str
100
+ words: List[WordAnalysis]
101
+ metadata: Dict
102
+
103
+ class QuranicDataset(Dataset):
104
+ def __init__(self, processed_data: List[Dict], tokenizer):
105
+ self.examples = []
106
+ self.tokenizer = tokenizer
107
+ for verse_data in processed_data:
108
+ self.examples.extend(self._create_training_examples(verse_data))
109
+
110
+ def _create_training_examples(self, verse_data: Dict) -> List[Dict]:
111
+ examples = []
112
+ text_block = (
113
+ f"[VERSE {verse_data['chapter']}:{verse_data['verse']}]\n"
114
+ f"Arabic: {verse_data['arabic_text']}\n"
115
+ f"Translation: {verse_data['translation']}\n"
116
+ "Morphological Analysis:\n"
117
+ )
118
+ for word in verse_data['words']:
119
+ text_block += (
120
+ f"[WORD] {word['arabic']}\n"
121
+ f"Root: {word['root']}\n"
122
+ f"Features: {', '.join(word['features'])}\n"
123
+ )
124
+ examples.append(self._format_example(text_block))
125
+ return examples
126
+
127
+ def _format_example(self, text: str) -> Dict:
128
+ encodings = self.tokenizer(
129
+ text,
130
+ truncation=True,
131
+ max_length=64,
132
+ padding="max_length",
133
+ return_tensors="pt"
134
+ )
135
+ return {
136
+ "input_ids": encodings["input_ids"][0].cpu(),
137
+ "attention_mask": encodings["attention_mask"][0].cpu()
138
+ }
139
+
140
+ def __len__(self):
141
+ return len(self.examples)
142
+
143
+ def __getitem__(self, idx):
144
+ return self.examples[idx]
145
+
146
+ class QuranicDataProcessor:
147
+ def __init__(self, source_dir: str, output_dir: str):
148
+ self.source_dir = source_dir
149
+ self.output_dir = output_dir
150
+ self.morphological_data: Dict[str, Dict] = {}
151
+ self.word_by_word_data: Dict[str, List[str]] = {}
152
+ self.translation_data: Dict[str, str] = {}
153
+ self.processing_lock = Lock()
154
+ os.makedirs(output_dir, exist_ok=True)
155
+ os.makedirs(os.path.join(output_dir, 'json'), exist_ok=True)
156
+ os.makedirs(os.path.join(output_dir, 'txt'), exist_ok=True)
157
+ os.makedirs(os.path.join(output_dir, 'checkpoints'), exist_ok=True)
158
+ logger.info(f"Initialized processor with source dir: {source_dir}")
159
+
160
+ def load_source_files(self) -> bool:
161
+ try:
162
+ logger.info("Loading morphological data...")
163
+ morph_path = os.path.join(self.source_dir, 'quranic-corpus-morphology-0.4.txt')
164
+ with open(morph_path, 'r', encoding='utf-8') as f:
165
+ next(f)
166
+ for line in f:
167
+ if line.strip() and not line.startswith('#'):
168
+ parts = line.strip().split('\t')
169
+ if len(parts) >= 4:
170
+ location = parts[0].strip('()')
171
+ self.morphological_data[location] = {
172
+ 'form': parts[1],
173
+ 'tag': parts[2],
174
+ 'features': parts[3]
175
+ }
176
+ logger.info(f"Loaded {len(self.morphological_data)} morphological entries")
177
+ logger.info("Loading translation data...")
178
+ trans_path = os.path.join(self.source_dir, 'en.sample.quran-maududi.txt')
179
+ with open(trans_path, 'r', encoding='utf-8') as f:
180
+ next(f)
181
+ for line in f:
182
+ if line.strip():
183
+ parts = line.strip().split('|')
184
+ if len(parts) >= 3:
185
+ key = f"{parts[0]}:{parts[1]}"
186
+ self.translation_data[key] = parts[2].strip()
187
+ logger.info(f"Loaded {len(self.translation_data)} verse translations")
188
+ logger.info("Loading word-by-word data...")
189
+ word_path = os.path.join(self.source_dir, 'en.w4w.qurandev.txt')
190
+ with open(word_path, 'r', encoding='utf-8-sig') as f:
191
+ lines = [line.strip() for line in f if line.strip()]
192
+ sorted_keys = sorted(self.translation_data.keys(), key=lambda x: (int(x.split(':')[0]), int(x.split(':')[1])))
193
+ if len(lines) != len(sorted_keys):
194
+ logger.warning("Mismatch between word-by-word file and translation data")
195
+ for i, verse_key in enumerate(sorted_keys):
196
+ if i < len(lines):
197
+ words = [w.strip() for w in lines[i].split('|') if w.strip()]
198
+ self.word_by_word_data[verse_key] = words
199
+ logger.info(f"Loaded word-by-word data for {len(self.word_by_word_data)} verses")
200
+ return True
201
+ except Exception as e:
202
+ logger.error(f"Error loading source files: {str(e)}")
203
+ logger.error(traceback.format_exc())
204
+ return False
205
+
206
+ def process_verse(self, chapter: int, verse: int) -> Optional[VerseData]:
207
+ try:
208
+ verse_ref = f"{chapter}:{verse}"
209
+ logger.info(f"Processing verse {verse_ref}")
210
+ translation = self.translation_data.get(verse_ref)
211
+ if not translation:
212
+ logger.warning(f"No translation for verse {verse_ref}")
213
+ return None
214
+ verse_word_list = self.word_by_word_data.get(verse_ref, [])
215
+ if not verse_word_list:
216
+ logger.warning(f"No word-by-word data for verse {verse_ref}")
217
+ return None
218
+ verse_words: List[WordAnalysis] = []
219
+ arabic_text = ""
220
+ for pos in range(1, len(verse_word_list) + 1):
221
+ pattern = f"{chapter}:{verse}:{pos}:"
222
+ matching_entries = [data for loc, data in self.morphological_data.items() if loc.startswith(pattern)]
223
+ if not matching_entries:
224
+ logger.debug(f"No morphological data for {pattern}")
225
+ continue
226
+ combined_form = " ".join(entry['form'] for entry in matching_entries)
227
+ combined_features = []
228
+ root = ""
229
+ for entry in matching_entries:
230
+ features = entry['features'].split('|')
231
+ combined_features.extend(features)
232
+ if not root:
233
+ for f in features:
234
+ if 'ROOT:' in f:
235
+ root = f.split('ROOT:')[1]
236
+ break
237
+ word_translation = verse_word_list[pos - 1]
238
+ word = WordAnalysis(
239
+ arabic=combined_form,
240
+ translation=word_translation,
241
+ position=str(pos),
242
+ morphology=matching_entries[0],
243
+ features=combined_features,
244
+ root=root,
245
+ location=f"{chapter}:{verse}:{pos}",
246
+ metadata={}
247
+ )
248
+ verse_words.append(word)
249
+ arabic_text += f" {combined_form}"
250
+ verse_data = VerseData(
251
+ chapter=chapter,
252
+ verse=verse,
253
+ arabic_text=arabic_text.strip(),
254
+ translation=translation,
255
+ words=verse_words,
256
+ metadata={
257
+ "processed_timestamp": datetime.now().isoformat(),
258
+ "word_count": len(verse_words)
259
+ }
260
+ )
261
+ self._save_verse_data(verse_data)
262
+ return verse_data
263
+ except Exception as e:
264
+ logger.error(f"Error processing verse {chapter}:{verse}: {str(e)}")
265
+ logger.error(traceback.format_exc())
266
+ return None
267
+
268
+ def _save_verse_data(self, verse_data: VerseData):
269
+ try:
270
+ verse_ref = f"{verse_data.chapter}:{verse_data.verse}"
271
+ json_path = os.path.join(self.output_dir, 'json', f'verse_{verse_ref.replace(":", "_")}.json')
272
+ with open(json_path, 'w', encoding='utf-8') as f:
273
+ json.dump(asdict(verse_data), f, ensure_ascii=False, indent=2)
274
+ txt_path = os.path.join(self.output_dir, 'txt', f'verse_{verse_ref.replace(":", "_")}.txt')
275
+ with open(txt_path, 'w', encoding='utf-8') as f:
276
+ f.write(f"=== Verse {verse_ref} ===\n\n")
277
+ f.write(f"Arabic Text:\n{verse_data.arabic_text}\n\n")
278
+ f.write(f"Translation:\n{verse_data.translation}\n\n")
279
+ f.write("Word Analysis:\n")
280
+ for i, word in enumerate(verse_data.words, 1):
281
+ f.write(f"\nWord {i}:\n")
282
+ f.write(f" Arabic: {word.arabic}\n")
283
+ f.write(f" Translation: {word.translation}\n")
284
+ f.write(f" Root: {word.root}\n")
285
+ f.write(" Features:\n")
286
+ for feature in word.features:
287
+ f.write(f" - {feature}\n")
288
+ f.write("\n")
289
+ logger.info(f"Saved verse data to {json_path} and {txt_path}")
290
+ except Exception as e:
291
+ logger.error(f"Error saving verse data: {str(e)}")
292
+ logger.error(traceback.format_exc())
293
+
294
+ class QuranicModelTrainer:
295
+ def __init__(self,
296
+ model_name: str = "google/gemma-2-2b",
297
+ processed_data_dir: str = "processed_data",
298
+ checkpoint_dir: str = "checkpoints"):
299
+ self.processed_data_dir = processed_data_dir
300
+ self.checkpoint_dir = checkpoint_dir
301
+ self.accelerator = Accelerator()
302
+ logger.info("Initializing Accelerator...")
303
+
304
+ self.tokenizer = AutoTokenizer.from_pretrained(
305
+ model_name,
306
+ token=os.environ.get("HF_TOKEN"),
307
+ additional_special_tokens=["[VERSE]", "[WORD]", "[ROOT]", "[FEATURES]"],
308
+ trust_remote_code=True
309
+ )
310
+ if self.tokenizer.pad_token is None:
311
+ self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
312
+
313
+ try:
314
+ self.model = AutoModelForCausalLM.from_pretrained(
315
+ model_name,
316
+ token=os.environ.get("HF_TOKEN"),
317
+ torch_dtype=torch.float32,
318
+ low_cpu_mem_usage=True,
319
+ trust_remote_code=True,
320
+ attn_implementation="eager"
321
+ )
322
+ except Exception as e:
323
+ logger.error(f"Error loading model directly: {str(e)}")
324
+ logger.info("Attempting to load with fallback parameters...")
325
+ from transformers import AutoConfig
326
+ config = AutoConfig.from_pretrained(
327
+ model_name,
328
+ token=os.environ.get("HF_TOKEN"),
329
+ trust_remote_code=True
330
+ )
331
+ self.model = AutoModelForCausalLM.from_pretrained(
332
+ model_name,
333
+ token=os.environ.get("HF_TOKEN"),
334
+ config=config,
335
+ torch_dtype=torch.float32,
336
+ low_cpu_mem_usage=True,
337
+ trust_remote_code=True,
338
+ revision="main",
339
+ attn_implementation="eager"
340
+ )
341
+
342
+ self.model.resize_token_embeddings(len(self.tokenizer))
343
+ self.model.train()
344
+ self.model.config.use_cache = False
345
+ self.model = self.accelerator.prepare(self.model)
346
+
347
+ if hasattr(self.model, "gradient_checkpointing_enable"):
348
+ self.model.gradient_checkpointing_enable()
349
+ else:
350
+ logger.warning("Gradient checkpointing not available for this model")
351
+
352
+ def prepare_training_data(self, chapter_data: List[Dict]) -> Dataset:
353
+ return QuranicDataset(chapter_data, self.tokenizer)
354
+
355
+ def train_chunk(self, training_args: TrainingArguments, dataset: Dataset, chunk_output_dir: str) -> bool:
356
+ try:
357
+ data_collator = DataCollatorForLanguageModeling(
358
+ tokenizer=self.tokenizer,
359
+ mlm=False
360
+ )
361
+ trainer = Trainer(
362
+ model=self.model,
363
+ args=training_args,
364
+ train_dataset=dataset,
365
+ tokenizer=self.tokenizer,
366
+ data_collator=data_collator
367
+ )
368
+ logger.info(f"Starting training on chunk at {chunk_output_dir} with device {self.accelerator.device}")
369
+ trainer.train()
370
+ trainer.save_model(chunk_output_dir)
371
+ zip_filename = zip_checkpoint(chunk_output_dir)
372
+ base_url = os.environ.get("HF_SPACE_URL", "http://localhost")
373
+ download_link = f"{base_url}/file/{zip_filename}"
374
+ logger.info(f"Checkpoint download link: {download_link}")
375
+ with open(os.path.join(chunk_output_dir, "download_link.txt"), "w") as f:
376
+ f.write(download_link)
377
+ del trainer
378
+ gc.collect()
379
+ manage_memory()
380
+ manage_gpu_resources()
381
+ return True
382
+ except RuntimeError as e:
383
+ if "NVML_SUCCESS" in str(e):
384
+ logger.error(f"Error in training chunk at {chunk_output_dir}: {str(e)}")
385
+ logger.info("GPU error detected. Shifting to CPU...")
386
+ if torch.cuda.is_available():
387
+ torch.cuda.empty_cache()
388
+ self.model = self.model.to("cpu")
389
+ training_args.no_cuda = True
390
+ try:
391
+ trainer = Trainer(
392
+ model=self.model,
393
+ args=training_args,
394
+ train_dataset=dataset,
395
+ tokenizer=self.tokenizer,
396
+ data_collator=data_collator
397
+ )
398
+ logger.info(f"Retrying training on CPU for chunk at {chunk_output_dir}")
399
+ trainer.train()
400
+ trainer.save_model(chunk_output_dir)
401
+ zip_filename = zip_checkpoint(chunk_output_dir)
402
+ base_url = os.environ.get("HF_SPACE_URL", "http://localhost")
403
+ download_link = f"{base_url}/file/{zip_filename}"
404
+ logger.info(f"Checkpoint download link: {download_link}")
405
+ with open(os.path.join(chunk_output_dir, "download_link.txt"), "w") as f:
406
+ f.write(download_link)
407
+ del trainer
408
+ gc.collect()
409
+ manage_memory()
410
+ return True
411
+ except Exception as cpu_e:
412
+ logger.error(f"Training failed on CPU: {str(cpu_e)}")
413
+ logger.error(traceback.format_exc())
414
+ return False
415
+ else:
416
+ logger.error(f"Error in training chunk at {chunk_output_dir}: {str(e)}")
417
+ logger.error(traceback.format_exc())
418
+ return False
419
+
420
+ def poll_for_gpu(self, poll_interval: int = 10, max_attempts: int = 30) -> bool:
421
+ attempts = 0
422
+ while attempts < max_attempts:
423
+ if torch.cuda.is_available():
424
+ manage_gpu_resources(1)
425
+ logger.info("GPU is now available for training.")
426
+ return True
427
+ time.sleep(poll_interval)
428
+ attempts += 1
429
+ logger.info(f"Polling for GPU availability... attempt {attempts}/{max_attempts}")
430
+ return False
431
+
432
+ def train_chapter(self,
433
+ chapter_num: int,
434
+ processed_verses: List[Dict],
435
+ chunk_size: int = 5,
436
+ num_train_epochs: int = 5,
437
+ per_device_train_batch_size: int = 1,
438
+ learning_rate: float = 3e-5,
439
+ weight_decay: float = 0.01,
440
+ gradient_accumulation_steps: int = 32) -> bool:
441
+ total_examples = len(processed_verses)
442
+ total_chunks = math.ceil(total_examples / chunk_size)
443
+ logger.info(f"Chapter {chapter_num}: {total_examples} examples, {total_chunks} chunks.")
444
+ for chunk_index in range(total_chunks):
445
+ chunk_data = processed_verses[chunk_index * chunk_size: (chunk_index + 1) * chunk_size]
446
+ dataset = self.prepare_training_data(chunk_data)
447
+ chunk_output_dir = os.path.join(self.checkpoint_dir, f"chapter_{chapter_num}", f"chunk_{chunk_index}")
448
+ os.makedirs(chunk_output_dir, exist_ok=True)
449
+
450
+ # For T4, enable FP16 training for better performance.
451
+ training_args = TrainingArguments(
452
+ output_dir=chunk_output_dir,
453
+ overwrite_output_dir=True,
454
+ num_train_epochs=num_train_epochs,
455
+ per_device_train_batch_size=per_device_train_batch_size,
456
+ learning_rate=learning_rate,
457
+ weight_decay=weight_decay,
458
+ gradient_accumulation_steps=gradient_accumulation_steps,
459
+ fp16=(self.accelerator.device.type == "cuda"),
460
+ remove_unused_columns=False,
461
+ logging_steps=50,
462
+ report_to="none",
463
+ eval_strategy="no",
464
+ no_cuda=(self.accelerator.device.type != "cuda"),
465
+ optim="adamw_torch",
466
+ dataloader_num_workers=0,
467
+ dataloader_pin_memory=False
468
+ )
469
+ logger.info(f"Training chunk {chunk_index+1}/{total_chunks} for Chapter {chapter_num}...")
470
+ success = self.train_chunk(training_args, dataset, chunk_output_dir)
471
+
472
+ if not success and self.accelerator.device.type == "cuda":
473
+ logger.error(f"Training failed for Chapter {chapter_num} on chunk {chunk_index+1}.")
474
+ return False
475
+ logger.info(f"Completed training for Chapter {chapter_num}")
476
+ return True
477
+
478
+ class QuranicPipeline:
479
+ def __init__(self,
480
+ source_dir: str = ".",
481
+ working_dir: str = "working_directory",
482
+ start_chapter: int = 1,
483
+ end_chapter: int = 114):
484
+ self.source_dir = source_dir
485
+ self.working_dir = working_dir
486
+ self.start_chapter = start_chapter
487
+ self.end_chapter = end_chapter
488
+ self.setup_directories()
489
+ global logger
490
+ logger = logging.getLogger(__name__)
491
+ self.state = {
492
+ "last_processed_chapter": 0,
493
+ "last_trained_chapter": 0,
494
+ "current_state": "initialized",
495
+ "errors": [],
496
+ "start_time": datetime.now().isoformat()
497
+ }
498
+ self.load_state()
499
+ try:
500
+ logger.info("Initializing Quranic Data Processor...")
501
+ self.processor = QuranicDataProcessor(
502
+ source_dir=self.source_dir,
503
+ output_dir=os.path.join(self.working_dir, "processed_data")
504
+ )
505
+ logger.info("Initializing Quranic Model Trainer...")
506
+ self.trainer = QuranicModelTrainer(
507
+ model_name="google/gemma-2-2b",
508
+ processed_data_dir=os.path.join(self.working_dir, "processed_data"),
509
+ checkpoint_dir=os.path.join(self.working_dir, "checkpoints")
510
+ )
511
+ self.state["current_state"] = "ready"
512
+ self.save_state()
513
+ except Exception as e:
514
+ self.handle_error("Initialization failed", e)
515
+ raise
516
+
517
+ def setup_directories(self):
518
+ dirs = [
519
+ self.working_dir,
520
+ os.path.join(self.working_dir, "processed_data"),
521
+ os.path.join(self.working_dir, "checkpoints"),
522
+ os.path.join(self.working_dir, "logs"),
523
+ os.path.join(self.working_dir, "state")
524
+ ]
525
+ for d in dirs:
526
+ os.makedirs(d, exist_ok=True)
527
+
528
+ def load_state(self):
529
+ state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
530
+ if os.path.exists(state_file):
531
+ try:
532
+ with open(state_file, 'r') as f:
533
+ saved_state = json.load(f)
534
+ self.state.update(saved_state)
535
+ logger.info(f"Loaded previous state: Last processed chapter {self.state.get('last_processed_chapter')}, last trained chapter {self.state.get('last_trained_chapter')}")
536
+ except Exception as e:
537
+ logger.warning(f"Could not load previous state: {str(e)}")
538
+
539
+ def save_state(self):
540
+ state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
541
+ with open(state_file, 'w') as f:
542
+ json.dump(self.state, f, indent=2)
543
+
544
+ def handle_error(self, context: str, error: Exception):
545
+ error_detail = {
546
+ "timestamp": datetime.now().isoformat(),
547
+ "context": context,
548
+ "error": str(error),
549
+ "traceback": traceback.format_exc()
550
+ }
551
+ self.state.setdefault("errors", []).append(error_detail)
552
+ logger.error(f"{context}: {str(error)}")
553
+ self.save_state()
554
+
555
+ def run_pipeline(self):
556
+ logger.info("Starting pipeline execution")
557
+ try:
558
+ if not self.processor.load_source_files():
559
+ raise Exception("Failed to load source files")
560
+ for chapter in range(self.start_chapter, self.end_chapter + 1):
561
+ logger.info(f"=== Processing Chapter {chapter} ===")
562
+ processed_chapter_data = []
563
+ verse = 1
564
+ while True:
565
+ verse_data = self.processor.process_verse(chapter, verse)
566
+ if verse_data is None:
567
+ break
568
+ processed_chapter_data.append(asdict(verse_data))
569
+ verse += 1
570
+ if processed_chapter_data:
571
+ success = self.trainer.train_chapter(chapter, processed_chapter_data)
572
+ if not success:
573
+ logger.error(f"Training failed for Chapter {chapter}. Stopping pipeline.")
574
+ break
575
+ self.state["last_trained_chapter"] = chapter
576
+ self.save_state()
577
+ else:
578
+ logger.warning(f"No processed data for Chapter {chapter}")
579
+ self.state["last_processed_chapter"] = chapter
580
+ self.save_state()
581
+ manage_memory()
582
+ manage_gpu_resources()
583
+ logger.info("Pipeline execution completed")
584
+ final_model_dir = os.path.join(self.working_dir, "final_model")
585
+ os.makedirs(final_model_dir, exist_ok=True)
586
+ self.trainer.accelerator.wait_for_everyone()
587
+ self.trainer.accelerator.save_model(self.trainer.model, final_model_dir)
588
+ self.trainer.tokenizer.save_pretrained(final_model_dir)
589
+ logger.info(f"Final model saved to {final_model_dir}")
590
+ except Exception as e:
591
+ self.handle_error("Pipeline execution failed", e)
592
+ raise
593
+
594
+ @spaces.GPU()
595
+ def start_pipeline():
596
+ try:
597
+ logger.info("Starting Quranic Training Pipeline with Gemma-2-2b on T4 Medium")
598
+ logger.info(f"PyTorch version: {torch.__version__}")
599
+ logger.info(f"CUDA available: {torch.cuda.is_available()}")
600
+ if torch.cuda.is_available():
601
+ logger.info(f"CUDA device count: {torch.cuda.device_count()}")
602
+ logger.info(f"CUDA device name: {torch.cuda.get_device_name(0)}")
603
+ if not os.environ.get("HF_TOKEN"):
604
+ logger.warning("HF_TOKEN environment variable not set. Model loading may fail.")
605
+ required_files = [
606
+ 'quranic-corpus-morphology-0.4.txt',
607
+ 'en.sample.quran-maududi.txt',
608
+ 'en.w4w.qurandev.txt'
609
+ ]
610
+ missing_files = [f for f in required_files if not os.path.exists(f)]
611
+ if missing_files:
612
+ return f"Missing required data files: {', '.join(missing_files)}"
613
+ pipeline = QuranicPipeline(
614
+ source_dir=".",
615
+ working_dir="working_directory",
616
+ start_chapter=1,
617
+ end_chapter=114
618
+ )
619
+ pipeline.run_pipeline()
620
+ return "Pipeline execution completed successfully."
621
+ except Exception as e:
622
+ error_msg = f"Pipeline execution failed: {str(e)}\n{traceback.format_exc()}"
623
+ logger.error(error_msg)
624
+ return error_msg
625
+
626
+ iface = gr.Interface(
627
+ fn=start_pipeline,
628
+ inputs=[],
629
+ outputs=gr.Textbox(label="Pipeline Status", lines=10),
630
+ title="Quranic Training Pipeline for Gemma-2-2b on T4 Medium",
631
+ description="""This pipeline is updated for T4 medium with FP16 training,
632
+ checkpoint saving (download link provided), and enhanced error handling.
633
+ """
634
+ )
635
+
636
+ if __name__ == "__main__":
637
+ iface.launch()