eBlessings commited on
Commit
0763938
·
verified ·
1 Parent(s): ab53392

Upload t4d.py

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