eBlessings commited on
Commit
d10f83f
·
verified ·
1 Parent(s): e0f89d5

Delete app.py

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