Update app.py

app.py

@@ -1,7 +1,527 @@
+#!/usr/bin/env python3
+"""
+app.py – Quranic Data Training Pipeline Endpoint for ZeroGPU Spaces
+--------------------------------------------------------------------
+This script integrates a full Quranic data processing and training pipeline
+into a Gradio interface endpoint. It is optimized for CPU-based training on 
+Hugging Face ZeroGPU (using the Gradio SDK) and uses chunked incremental 
+training, memory management, and gradient checkpointing to efficiently update 
+Google’s Gemma-2B model with Quranic data.
+
+Requirements:
+ - Hugging Face Transformers, Gradio, PyTorch, psutil
+ - Hugging Face PRO account with ZeroGPU enabled (set ZEROGPU_V2 and ZERO_GPU_PATCH_TORCH_DEVICE env variables)
+ - Ubuntu CPU/Linux with access to ZeroGPU hardware via Spaces
+ - Source data in "source_files" directory
+ - Sufficient storage in "working_directory"
+
+Author: [M-Saddam Hussain]
+Date: February 2025
+Data References: [Tanzil.net, IslamSource, QuranicCorpus]
+"""
+
+import json
+import logging
+import os
+import sys
+import traceback
+import gc
+import time
+import psutil
+import math
+from datetime import datetime
+from typing import Dict, List, Optional
+from dataclasses import dataclass, asdict
+
+import torch
+# Limit PyTorch threads for CPU stability.
+torch.set_num_threads(8)
+
+from torch.utils.data import Dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from threading import Lock
+
+# Import Gradio and the spaces module for ZeroGPU.
 import gradio as gr
+import spaces
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('pipeline.log'),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
+def manage_memory(threshold_percent: int = 90, min_available_mb: int = 500, sleep_duration: int = 10):
+    """
+    Check memory usage; if usage is high or available memory is low,
+    force garbage collection and sleep briefly.
+    """
+    vm = psutil.virtual_memory()
+    used_percent = vm.percent
+    available_mb = vm.available / (1024 * 1024)
+    logger.info(f"Memory usage: {used_percent}% used, {available_mb:.2f} MB available")
+    if used_percent > threshold_percent or available_mb < min_available_mb:
+        logger.warning("High memory usage detected, forcing garbage collection and sleeping...")
+        gc.collect()
+        time.sleep(sleep_duration)
+
+@dataclass
+class WordAnalysis:
+    """Structured representation of word-level analysis"""
+    arabic: str
+    translation: str
+    position: str
+    morphology: Dict
+    features: List[str]
+    root: str
+    location: str
+    metadata: Dict
+
+@dataclass
+class VerseData:
+    """Structured representation of verse-level data"""
+    chapter: int
+    verse: int
+    arabic_text: str
+    translation: str
+    words: List[WordAnalysis]
+    metadata: Dict
+
+class QuranicDataset(Dataset):
+    """Custom dataset for Quranic text training."""
+    def __init__(self, processed_data: List[Dict], tokenizer):
+        self.examples = []
+        self.tokenizer = tokenizer
+        for verse_data in processed_data:
+            self.examples.extend(self._create_training_examples(verse_data))
+    
+    def _create_training_examples(self, verse_data: Dict) -> List[Dict]:
+        examples = []
+        text_block = (
+            f"[VERSE {verse_data['chapter']}:{verse_data['verse']}]\n"
+            f"Arabic: {verse_data['arabic_text']}\n"
+            f"Translation: {verse_data['translation']}\n"
+            "Morphological Analysis:\n"
+        )
+        for word in verse_data['words']:
+            text_block += (
+                f"[WORD] {word['arabic']}\n"
+                f"Root: {word['root']}\n"
+                f"Features: {', '.join(word['features'])}\n"
+            )
+        examples.append(self._format_example(text_block))
+        return examples
+    
+    def _format_example(self, text: str) -> Dict:
+        encodings = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=64,  # Reduced length to lower memory usage.
+            padding="max_length",
+            return_tensors="pt"
+        )
+        return {
+            "input_ids": encodings["input_ids"][0],
+            "attention_mask": encodings["attention_mask"][0]
+        }
+    
+    def __len__(self):
+        return len(self.examples)
+    
+    def __getitem__(self, idx):
+        return self.examples[idx]
+
+class QuranicDataProcessor:
+    """Processes Quranic data into structured training examples."""
+    def __init__(self, source_dir: str, output_dir: str):
+        self.source_dir = source_dir
+        self.output_dir = output_dir
+        self.morphological_data: Dict[str, Dict] = {}
+        self.word_by_word_data: Dict[str, List[str]] = {}
+        self.translation_data: Dict[str, str] = {}
+        self.processed_verses = set()
+        self.processing_lock = Lock()
+        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs(os.path.join(output_dir, 'json'), exist_ok=True)
+        os.makedirs(os.path.join(output_dir, 'txt'), exist_ok=True)
+        os.makedirs(os.path.join(output_dir, 'checkpoints'), exist_ok=True)
+        logger.info(f"Initialized processor with source dir: {source_dir}")
+    
+    def load_source_files(self) -> bool:
+        """Loads morphological, translation, and word-by-word data."""
+        try:
+            logger.info("Loading morphological data...")
+            morph_path = os.path.join(self.source_dir, 'quranic-corpus-morphology-0.4.txt')
+            with open(morph_path, 'r', encoding='utf-8') as f:
+                next(f)
+                for line in f:
+                    if line.strip() and not line.startswith('#'):
+                        parts = line.strip().split('\t')
+                        if len(parts) >= 4:
+                            location = parts[0].strip('()')
+                            self.morphological_data[location] = {
+                                'form': parts[1],
+                                'tag': parts[2],
+                                'features': parts[3]
+                            }
+            logger.info(f"Loaded {len(self.morphological_data)} morphological entries")
+            logger.info("Loading translation data...")
+            trans_path = os.path.join(self.source_dir, 'en.sample.quran-maududi.txt')
+            with open(trans_path, 'r', encoding='utf-8') as f:
+                next(f)
+                for line in f:
+                    if line.strip():
+                        parts = line.strip().split('|')
+                        if len(parts) >= 3:
+                            key = f"{parts[0]}:{parts[1]}"
+                            self.translation_data[key] = parts[2].strip()
+            logger.info(f"Loaded {len(self.translation_data)} verse translations")
+            logger.info("Loading word-by-word data...")
+            word_path = os.path.join(self.source_dir, 'en.w4w.qurandev.txt')
+            with open(word_path, 'r', encoding='utf-8-sig') as f:
+                lines = [line.strip() for line in f if line.strip()]
+            sorted_keys = sorted(self.translation_data.keys(), key=lambda x: (int(x.split(':')[0]), int(x.split(':')[1])))
+            if len(lines) != len(sorted_keys):
+                logger.warning("Mismatch between word-by-word file and translation data")
+            for i, verse_key in enumerate(sorted_keys):
+                if i < len(lines):
+                    words = [w.strip() for w in lines[i].split('|') if w.strip()]
+                    self.word_by_word_data[verse_key] = words
+            logger.info(f"Loaded word-by-word data for {len(self.word_by_word_data)} verses")
+            return True
+        except Exception as e:
+            logger.error(f"Error loading source files: {str(e)}")
+            logger.error(traceback.format_exc())
+            return False
+
+    def process_verse(self, chapter: int, verse: int) -> Optional[VerseData]:
+        """Processes a single verse into structured format."""
+        try:
+            verse_ref = f"{chapter}:{verse}"
+            logger.info(f"Processing verse {verse_ref}")
+            translation = self.translation_data.get(verse_ref)
+            if not translation:
+                logger.warning(f"No translation for verse {verse_ref}")
+                return None
+            verse_word_list = self.word_by_word_data.get(verse_ref, [])
+            if not verse_word_list:
+                logger.warning(f"No word-by-word data for verse {verse_ref}")
+                return None
+            verse_words: List[WordAnalysis] = []
+            arabic_text = ""
+            for pos in range(1, len(verse_word_list) + 1):
+                pattern = f"{chapter}:{verse}:{pos}:"
+                matching_entries = [data for loc, data in self.morphological_data.items() if loc.startswith(pattern)]
+                if not matching_entries:
+                    logger.debug(f"No morphological data for {pattern}")
+                    continue
+                combined_form = " ".join(entry['form'] for entry in matching_entries)
+                combined_features = []
+                root = ""
+                for entry in matching_entries:
+                    features = entry['features'].split('|')
+                    combined_features.extend(features)
+                    if not root:
+                        for f in features:
+                            if 'ROOT:' in f:
+                                root = f.split('ROOT:')[1]
+                                break
+                word_translation = verse_word_list[pos - 1]
+                word = WordAnalysis(
+                    arabic=combined_form,
+                    translation=word_translation,
+                    position=str(pos),
+                    morphology=matching_entries[0],
+                    features=combined_features,
+                    root=root,
+                    location=f"{chapter}:{verse}:{pos}",
+                    metadata={}
+                )
+                verse_words.append(word)
+                arabic_text += f" {combined_form}"
+            verse_data = VerseData(
+                chapter=chapter,
+                verse=verse,
+                arabic_text=arabic_text.strip(),
+                translation=translation,
+                words=verse_words,
+                metadata={
+                    "processed_timestamp": datetime.now().isoformat(),
+                    "word_count": len(verse_words)
+                }
+            )
+            self._save_verse_data(verse_data)
+            return verse_data
+        except Exception as e:
+            logger.error(f"Error processing verse {chapter}:{verse}: {str(e)}")
+            logger.error(traceback.format_exc())
+            return None
+
+    def _save_verse_data(self, verse_data: VerseData):
+        """Saves processed verse data as JSON and TXT."""
+        try:
+            verse_ref = f"{verse_data.chapter}:{verse_data.verse}"
+            json_path = os.path.join(self.output_dir, 'json', f'verse_{verse_ref.replace(":", "_")}.json')
+            with open(json_path, 'w', encoding='utf-8') as f:
+                json.dump(asdict(verse_data), f, ensure_ascii=False, indent=2)
+            txt_path = os.path.join(self.output_dir, 'txt', f'verse_{verse_ref.replace(":", "_")}.txt')
+            with open(txt_path, 'w', encoding='utf-8') as f:
+                f.write(f"=== Verse {verse_ref} ===\n\n")
+                f.write(f"Arabic Text:\n{verse_data.arabic_text}\n\n")
+                f.write(f"Translation:\n{verse_data.translation}\n\n")
+                f.write("Word Analysis:\n")
+                for i, word in enumerate(verse_data.words, 1):
+                    f.write(f"\nWord {i}:\n")
+                    f.write(f"  Arabic: {word.arabic}\n")
+                    f.write(f"  Translation: {word.translation}\n")
+                    f.write(f"  Root: {word.root}\n")
+                    f.write("  Features:\n")
+                    for feature in word.features:
+                        f.write(f"    - {feature}\n")
+                    f.write("\n")
+            logger.info(f"Saved verse data to {json_path} and {txt_path}")
+        except Exception as e:
+            logger.error(f"Error saving verse data: {str(e)}")
+            logger.error(traceback.format_exc())
+
+class QuranicModelTrainer:
+    """Trains the Gemma-2B model on Quranic data using chunked incremental updates."""
+    def __init__(self,
+                 model_name: str = "google/gemma-2-2b",
+                 processed_data_dir: str = "processed_data",
+                 checkpoint_dir: str = "checkpoints"):
+        self.processed_data_dir = processed_data_dir
+        self.checkpoint_dir = checkpoint_dir
+        self.device = "cpu"  # Training on CPU; ZeroGPU will handle GPU access.
+        logger.info("Loading tokenizer and model...")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            additional_special_tokens=["[VERSE]", "[WORD]", "[ROOT]", "[FEATURES]"]
+        )
+        self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True
+        )
+        self.model.resize_token_embeddings(len(self.tokenizer))
+        self.model.train()
+        # Disable caching and enable gradient checkpointing for memory savings.
+        self.model.config.use_cache = False
+        self.model.gradient_checkpointing_enable()
+    
+    def prepare_training_data(self, chapter_data: List[Dict]) -> Dataset:
+        """Creates a QuranicDataset from processed chapter data."""
+        return QuranicDataset(chapter_data, self.tokenizer)
+    
+    def train_chapter(self,
+                      chapter_num: int,
+                      processed_verses: List[Dict],
+                      chunk_size: int = 10,
+                      num_train_epochs: int = 10,
+                      per_device_train_batch_size: int = 1,
+                      learning_rate: float = 3e-5,
+                      weight_decay: float = 0.01,
+                      gradient_accumulation_steps: int = 64) -> bool:
+        """
+        Splits chapter data into chunks and trains incrementally to reduce memory usage.
+        """
+        try:
+            total_examples = len(processed_verses)
+            total_chunks = math.ceil(total_examples / chunk_size)
+            logger.info(f"Chapter {chapter_num}: {total_examples} examples, {total_chunks} chunks.")
+            for chunk_index in range(total_chunks):
+                chunk_data = processed_verses[chunk_index * chunk_size: (chunk_index + 1) * chunk_size]
+                dataset = self.prepare_training_data(chunk_data)
+                chunk_output_dir = os.path.join(self.checkpoint_dir, f"chapter_{chapter_num}", f"chunk_{chunk_index}")
+                os.makedirs(chunk_output_dir, exist_ok=True)
+                training_args = TrainingArguments(
+                    output_dir=chunk_output_dir,
+                    overwrite_output_dir=True,
+                    num_train_epochs=num_train_epochs,
+                    per_device_train_batch_size=per_device_train_batch_size,
+                    learning_rate=learning_rate,
+                    weight_decay=weight_decay,
+                    gradient_accumulation_steps=gradient_accumulation_steps,
+                    fp16=False,
+                    remove_unused_columns=False,
+                    logging_steps=50,
+                    report_to="none",
+                    evaluation_strategy="no",
+                    no_cuda=True,
+                    dataloader_num_workers=0,
+                    dataloader_pin_memory=False
+                )
+                data_collator = DataCollatorForLanguageModeling(
+                    tokenizer=self.tokenizer,
+                    mlm=False
+                )
+                trainer = Trainer(
+                    model=self.model,
+                    args=training_args,
+                    train_dataset=dataset,
+                    tokenizer=self.tokenizer,
+                    data_collator=data_collator
+                )
+                logger.info(f"Training chunk {chunk_index+1}/{total_chunks} for Chapter {chapter_num}...")
+                trainer.train()
+                trainer.save_model(chunk_output_dir)
+                del trainer, dataset
+                gc.collect()
+                manage_memory()
+            logger.info(f"Completed training for Chapter {chapter_num}")
+            return True
+        except Exception as e:
+            logger.error(f"Error training chapter {chapter_num}: {str(e)}")
+            logger.error(traceback.format_exc())
+            return False
+
+class QuranicPipeline:
+    """Integrates data processing and incremental model training for all chapters."""
+    def __init__(self,
+                 source_dir: str = "source_files",
+                 working_dir: str = "working_directory",
+                 start_chapter: int = 1,
+                 end_chapter: int = 114):
+        self.source_dir = source_dir
+        self.working_dir = working_dir
+        self.start_chapter = start_chapter
+        self.end_chapter = end_chapter
+        self.setup_directories()
+        global logger
+        logger = logging.getLogger(__name__)
+        self.state = {
+            "last_processed_chapter": 0,
+            "last_trained_chapter": 0,
+            "current_state": "initialized",
+            "errors": [],
+            "start_time": datetime.now().isoformat()
+        }
+        self.load_state()
+        try:
+            logger.info("Initializing Quranic Data Processor...")
+            self.processor = QuranicDataProcessor(
+                source_dir=self.source_dir,
+                output_dir=os.path.join(self.working_dir, "processed_data")
+            )
+            logger.info("Initializing Quranic Model Trainer...")
+            self.trainer = QuranicModelTrainer(
+                model_name="google/gemma-2-2b",
+                processed_data_dir=os.path.join(self.working_dir, "processed_data"),
+                checkpoint_dir=os.path.join(self.working_dir, "checkpoints")
+            )
+            self.state["current_state"] = "ready"
+            self.save_state()
+        except Exception as e:
+            self.handle_error("Initialization failed", e)
+            raise
+
+    def setup_directories(self):
+        dirs = [
+            self.working_dir,
+            os.path.join(self.working_dir, "processed_data"),
+            os.path.join(self.working_dir, "checkpoints"),
+            os.path.join(self.working_dir, "logs"),
+            os.path.join(self.working_dir, "state")
+        ]
+        for d in dirs:
+            os.makedirs(d, exist_ok=True)
+
+    def load_state(self):
+        state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
+        if os.path.exists(state_file):
+            try:
+                with open(state_file, 'r') as f:
+                    saved_state = json.load(f)
+                self.state.update(saved_state)
+                logger.info(f"Loaded previous state: Last processed chapter {self.state.get('last_processed_chapter')}, "
+                            f"last trained chapter {self.state.get('last_trained_chapter')}")
+            except Exception as e:
+                logger.warning(f"Could not load previous state: {str(e)}")
+
+    def save_state(self):
+        state_file = os.path.join(self.working_dir, "state", "pipeline_state.json")
+        with open(state_file, 'w') as f:
+            json.dump(self.state, f, indent=2)
+
+    def handle_error(self, context: str, error: Exception):
+        error_detail = {
+            "timestamp": datetime.now().isoformat(),
+            "context": context,
+            "error": str(error),
+            "traceback": traceback.format_exc()
+        }
+        self.state.setdefault("errors", []).append(error_detail)
+        logger.error(f"{context}: {str(error)}")
+        self.save_state()
+
+    def run_pipeline(self):
+        """Runs processing and training for chapters sequentially."""
+        logger.info("Starting pipeline execution")
+        try:
+            if not self.processor.load_source_files():
+                raise Exception("Failed to load source files")
+            for chapter in range(self.start_chapter, self.end_chapter + 1):
+                logger.info(f"=== Processing Chapter {chapter} ===")
+                processed_chapter_data = []
+                verse = 1
+                while True:
+                    verse_data = self.processor.process_verse(chapter, verse)
+                    if verse_data is None:
+                        break
+                    processed_chapter_data.append(asdict(verse_data))
+                    verse += 1
+                if processed_chapter_data:
+                    success = self.trainer.train_chapter(chapter, processed_chapter_data)
+                    if not success:
+                        logger.error(f"Training failed for Chapter {chapter}. Stopping pipeline.")
+                        break
+                    self.state["last_trained_chapter"] = chapter
+                    self.save_state()
+                else:
+                    logger.warning(f"No processed data for Chapter {chapter}")
+                self.state["last_processed_chapter"] = chapter
+                self.save_state()
+                manage_memory()
+            logger.info("Pipeline execution completed")
+        except Exception as e:
+            self.handle_error("Pipeline execution failed", e)
+            raise
+
+# Define a Gradio endpoint function that triggers the training pipeline.
+@spaces.GPU()  # This decorator requests ZeroGPU hardware.
+def start_pipeline():
+    try:
+        pipeline = QuranicPipeline(
+            source_dir="source_files",
+            working_dir="working_directory",
+            start_chapter=1,
+            end_chapter=114
+        )
+        pipeline.run_pipeline()
+        return "Pipeline execution completed successfully."
+    except Exception as e:
+        return f"Pipeline execution failed: {str(e)}"
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Create a Gradio Interface with no inputs and a text output.
+iface = gr.Interface(
+    fn=start_pipeline,
+    inputs=[],
+    outputs="text",
+    title="Quranic Training Pipeline",
+    description="Click 'Submit' to trigger the Quranic data processing and training pipeline on ZeroGPU."
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    iface.launch()