Spaces:
Sleeping
Sleeping
Upload T4a.py
Browse files
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()
|