eBlessings commited on
Commit
ae45e4b
·
verified ·
1 Parent(s): 3c36c97

Upload gpt9.py

Browse files
Files changed (1) hide show
  1. gpt9.py +646 -0
gpt9.py ADDED
@@ -0,0 +1,646 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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_chapter(self,
396
+ chapter_num: int,
397
+ processed_verses: List[Dict],
398
+ chunk_size: int = 5, # Reduced chunk size to help with memory
399
+ num_train_epochs: int = 5, # Lower epochs for testing
400
+ per_device_train_batch_size: int = 1,
401
+ learning_rate: float = 3e-5,
402
+ weight_decay: float = 0.01,
403
+ gradient_accumulation_steps: int = 32) -> bool:
404
+ """
405
+ Splits chapter data into chunks and trains incrementally to reduce memory usage.
406
+ After each chunk, creates a downloadable zip of the checkpoint.
407
+ """
408
+ try:
409
+ total_examples = len(processed_verses)
410
+ total_chunks = math.ceil(total_examples / chunk_size)
411
+ logger.info(f"Chapter {chapter_num}: {total_examples} examples, {total_chunks} chunks.")
412
+ for chunk_index in range(total_chunks):
413
+ chunk_data = processed_verses[chunk_index * chunk_size: (chunk_index + 1) * chunk_size]
414
+ dataset = self.prepare_training_data(chunk_data)
415
+ chunk_output_dir = os.path.join(self.checkpoint_dir, f"chapter_{chapter_num}", f"chunk_{chunk_index}")
416
+ os.makedirs(chunk_output_dir, exist_ok=True)
417
+ # Set use_cpu dynamically based on GPU availability.
418
+ training_args = TrainingArguments(
419
+ output_dir=chunk_output_dir,
420
+ overwrite_output_dir=True,
421
+ num_train_epochs=num_train_epochs,
422
+ per_device_train_batch_size=per_device_train_batch_size,
423
+ learning_rate=learning_rate,
424
+ weight_decay=weight_decay,
425
+ gradient_accumulation_steps=gradient_accumulation_steps,
426
+ fp16=False,
427
+ remove_unused_columns=False,
428
+ logging_steps=50,
429
+ report_to="none",
430
+ # Updated per deprecation: use eval_strategy instead of evaluation_strategy.
431
+ eval_strategy="no",
432
+ use_cpu=not torch.cuda.is_available(),
433
+ dataloader_num_workers=0,
434
+ dataloader_pin_memory=False
435
+ )
436
+ data_collator = DataCollatorForLanguageModeling(
437
+ tokenizer=self.tokenizer,
438
+ mlm=False
439
+ )
440
+ trainer = Trainer(
441
+ model=self.model,
442
+ args=training_args,
443
+ train_dataset=dataset,
444
+ processing_class=self.tokenizer, # Updated per deprecation notice.
445
+ data_collator=data_collator
446
+ )
447
+ logger.info(f"Training chunk {chunk_index+1}/{total_chunks} for Chapter {chapter_num}...")
448
+ trainer.train()
449
+ trainer.save_model(chunk_output_dir)
450
+
451
+ # Zip the checkpoint folder and generate a download link.
452
+ zip_filename = zip_checkpoint(chunk_output_dir)
453
+ base_url = os.environ.get("HF_SPACE_URL", "http://localhost")
454
+ download_link = f"{base_url}/file/{zip_filename}"
455
+ logger.info(f"Checkpoint chunk {chunk_index+1} download link: {download_link}")
456
+ # Save the download link into a text file within the checkpoint directory.
457
+ with open(os.path.join(chunk_output_dir, "download_link.txt"), "w") as f:
458
+ f.write(download_link)
459
+
460
+ del trainer, dataset
461
+ gc.collect()
462
+ manage_memory()
463
+ manage_gpu_resources()
464
+ logger.info(f"Completed training for Chapter {chapter_num}")
465
+ return True
466
+ except Exception as e:
467
+ logger.error(f"Error training chapter {chapter_num}: {str(e)}")
468
+ logger.error(traceback.format_exc())
469
+ return False
470
+
471
+ class QuranicPipeline:
472
+ """Integrates data processing and incremental model training for all chapters."""
473
+ def __init__(self,
474
+ source_dir: str = ".",
475
+ working_dir: str = "working_directory",
476
+ start_chapter: int = 1,
477
+ end_chapter: int = 114):
478
+ self.source_dir = source_dir
479
+ self.working_dir = working_dir
480
+ self.start_chapter = start_chapter
481
+ self.end_chapter = end_chapter
482
+ self.setup_directories()
483
+ global logger
484
+ logger = logging.getLogger(__name__)
485
+ self.state = {
486
+ "last_processed_chapter": 0,
487
+ "last_trained_chapter": 0,
488
+ "current_state": "initialized",
489
+ "errors": [],
490
+ "start_time": datetime.now().isoformat()
491
+ }
492
+ self.load_state()
493
+ try:
494
+ logger.info("Initializing Quranic Data Processor...")
495
+ self.processor = QuranicDataProcessor(
496
+ source_dir=self.source_dir,
497
+ output_dir=os.path.join(self.working_dir, "processed_data")
498
+ )
499
+ logger.info("Initializing Quranic Model Trainer...")
500
+ self.trainer = QuranicModelTrainer(
501
+ model_name="google/gemma-2-2b",
502
+ processed_data_dir=os.path.join(self.working_dir, "processed_data"),
503
+ checkpoint_dir=os.path.join(self.working_dir, "checkpoints")
504
+ )
505
+ self.state["current_state"] = "ready"
506
+ self.save_state()
507
+ except Exception as e:
508
+ self.handle_error("Initialization failed", e)
509
+ raise
510
+
511
+ def setup_directories(self):
512
+ dirs = [
513
+ self.working_dir,
514
+ os.path.join(self.working_dir, "processed_data"),
515
+ os.path.join(self.working_dir, "checkpoints"),
516
+ os.path.join(self.working_dir, "logs"),
517
+ os.path.join(self.working_dir, "state")
518
+ ]
519
+ for d in dirs:
520
+ os.makedirs(d, exist_ok=True)
521
+
522
+ def load_state(self):
523
+ state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
524
+ if os.path.exists(state_file):
525
+ try:
526
+ with open(state_file, 'r') as f:
527
+ saved_state = json.load(f)
528
+ self.state.update(saved_state)
529
+ logger.info(f"Loaded previous state: Last processed chapter {self.state.get('last_processed_chapter')}, "
530
+ f"last trained chapter {self.state.get('last_trained_chapter')}")
531
+ except Exception as e:
532
+ logger.warning(f"Could not load previous state: {str(e)}")
533
+
534
+ def save_state(self):
535
+ state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
536
+ with open(state_file, 'w') as f:
537
+ json.dump(self.state, f, indent=2)
538
+
539
+ def handle_error(self, context: str, error: Exception):
540
+ error_detail = {
541
+ "timestamp": datetime.now().isoformat(),
542
+ "context": context,
543
+ "error": str(error),
544
+ "traceback": traceback.format_exc()
545
+ }
546
+ self.state.setdefault("errors", []).append(error_detail)
547
+ logger.error(f"{context}: {str(error)}")
548
+ self.save_state()
549
+
550
+ def run_pipeline(self):
551
+ """Runs processing and training for chapters sequentially, then saves the final model."""
552
+ logger.info("Starting pipeline execution")
553
+ try:
554
+ if not self.processor.load_source_files():
555
+ raise Exception("Failed to load source files")
556
+ for chapter in range(self.start_chapter, self.end_chapter + 1):
557
+ logger.info(f"=== Processing Chapter {chapter} ===")
558
+ processed_chapter_data = []
559
+ verse = 1
560
+ while True:
561
+ verse_data = self.processor.process_verse(chapter, verse)
562
+ if verse_data is None:
563
+ break
564
+ processed_chapter_data.append(asdict(verse_data))
565
+ verse += 1
566
+ if processed_chapter_data:
567
+ success = self.trainer.train_chapter(chapter, processed_chapter_data)
568
+ if not success:
569
+ logger.error(f"Training failed for Chapter {chapter}. Stopping pipeline.")
570
+ break
571
+ self.state["last_trained_chapter"] = chapter
572
+ self.save_state()
573
+ else:
574
+ logger.warning(f"No processed data for Chapter {chapter}")
575
+ self.state["last_processed_chapter"] = chapter
576
+ self.save_state()
577
+ manage_memory()
578
+ manage_gpu_resources()
579
+ logger.info("Pipeline execution completed")
580
+ # Save the final model and tokenizer after all training is complete.
581
+ final_model_dir = os.path.join(self.working_dir, "final_model")
582
+ os.makedirs(final_model_dir, exist_ok=True)
583
+ self.trainer.model.save_pretrained(final_model_dir)
584
+ self.trainer.tokenizer.save_pretrained(final_model_dir)
585
+ logger.info(f"Final model saved to {final_model_dir}")
586
+ except Exception as e:
587
+ self.handle_error("Pipeline execution failed", e)
588
+ raise
589
+
590
+ @spaces.GPU() # Request ZeroGPU hardware for the Space
591
+ def start_pipeline():
592
+ try:
593
+ logger.info("Starting Quranic Training Pipeline with Gemma-2-2b")
594
+ logger.info(f"PyTorch version: {torch.__version__}")
595
+ logger.info(f"CUDA available: {torch.cuda.is_available()}")
596
+ if torch.cuda.is_available():
597
+ logger.info(f"CUDA device count: {torch.cuda.device_count()}")
598
+ logger.info(f"CUDA device name: {torch.cuda.get_device_name(0)}")
599
+
600
+ if not os.environ.get("HF_TOKEN"):
601
+ logger.warning("HF_TOKEN environment variable not set. Model loading may fail.")
602
+
603
+ required_files = [
604
+ 'quranic-corpus-morphology-0.4.txt',
605
+ 'en.sample.quran-maududi.txt',
606
+ 'en.w4w.qurandev.txt'
607
+ ]
608
+ missing_files = [f for f in required_files if not os.path.exists(f)]
609
+ if missing_files:
610
+ return f"Missing required data files: {', '.join(missing_files)}"
611
+
612
+ pipeline = QuranicPipeline(
613
+ source_dir=".",
614
+ working_dir="working_directory",
615
+ start_chapter=1,
616
+ end_chapter=114
617
+ )
618
+ pipeline.run_pipeline()
619
+ return "Pipeline execution completed successfully."
620
+ except Exception as e:
621
+ error_msg = f"Pipeline execution failed: {str(e)}\n{traceback.format_exc()}"
622
+ logger.error(error_msg)
623
+ return error_msg
624
+
625
+ iface = gr.Interface(
626
+ fn=start_pipeline,
627
+ inputs=[],
628
+ outputs=gr.Textbox(label="Pipeline Status", lines=10),
629
+ title="Quranic Training Pipeline for Gemma-2-2b",
630
+ description="""This pipeline fine-tunes Google's Gemma-2-2b model on Quranic data.
631
+
632
+ Click 'Submit' to trigger the Quranic data processing and training pipeline on ZeroGPU.
633
+
634
+ Requirements:
635
+ - Transformers (>=4.42.0)
636
+ - Gradio (>=5.12.0)
637
+ - PyTorch (==2.2.2)
638
+ - psutil (==5.9.5)
639
+ - Accelerate (>=0.26.0)
640
+
641
+ The pipeline processes all 114 chapters of the Quran sequentially, with memory and GPU resource management optimizations for dynamic ZeroGPU environments.
642
+ Checkpoint download links are provided after every training chunk."""
643
+ )
644
+
645
+ if __name__ == "__main__":
646
+ iface.launch()