app.py

import gradio as gr
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForTokenClassification
from typing import List, Dict, Optional
import json

# ============================================================================
# CONFIGURATION
# ============================================================================
MODEL_NAME = "Badhon/Bangla-Punc-Restore-Model"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_LENGTH = 256  # Model's max length
CHUNK_SIZE = 200  # Process in chunks smaller than max length
OVERLAP = 20      # Overlap between chunks to maintain context

# ============================================================================
# LABEL MAPPING
# ============================================================================
id2punct = {
    0: "",        # O (no punctuation)
    1: ",",       # COMMA
    2: "।",       # DARI
    3: "?",       # QUESTION
    4: "!",       # EXCLAMATION
    5: ";",       # SEMICOLON
    6: ":",       # COLON
    7: "-"        # HYPHEN
}

label2id = {
    "O": 0,
    "COMMA": 1,
    "DARI": 2,
    "QUESTION": 3,
    "EXCLAMATION": 4,
    "SEMICOLON": 5,
    "COLON": 6,
    "HYPHEN": 7
}
id2label = {v: k for k, v in label2id.items()}


# ============================================================================
# LOAD MODEL
# ============================================================================
print(f"Loading model: {MODEL_NAME}")
print(f"Device: {DEVICE}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)
model.to(DEVICE)
model.eval()

print("Model loaded successfully!")


# ============================================================================
# INFERENCE FUNCTIONS
# ============================================================================
def predict_chunk(words: List[str]):
    """
    Predict punctuation for a chunk of words.
    
    Args:
        words: List of words to process
        
    Returns:
        List of predicted labels and confidences
    """
    # Tokenize
    encoding = tokenizer(
        words,
        is_split_into_words=True,
        truncation=True,
        max_length=MAX_LENGTH,
        return_tensors="pt"
    )
    
    # Move to device
    input_ids = encoding['input_ids'].to(DEVICE)
    attention_mask = encoding['attention_mask'].to(DEVICE)
    
    # Predict
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.logits
        predictions = torch.argmax(logits, dim=2)
    
    # Get word IDs
    word_ids = encoding.word_ids(batch_index=0)
    
    # Align predictions with words
    word_predictions = []
    word_confidences = []
    prev_word_id = None
    
    for idx, word_id in enumerate(word_ids):
        if word_id is not None and word_id != prev_word_id:
            pred_label = predictions[0][idx].item()
            word_predictions.append(pred_label)
            
            # Calculate confidence
            probs = torch.softmax(logits[0][idx], dim=0)
            confidence = probs[pred_label].item()
            word_confidences.append(confidence)
            
            prev_word_id = word_id
    
    return word_predictions, word_confidences


def predict_punctuation(text: str, show_details: bool = False):
    """
    Restore punctuation for input text (handles long texts via chunking).
    
    Args:
        text: Input text without punctuation
        show_details: Whether to show detailed information
        
    Returns:
        Restored text and optional detailed information
    """
    # Handle empty input
    if not text or not text.strip():
        return "⚠️ Please enter some text!", ""
    
    # Split into words
    words = text.split()
    if not words:
        return text, ""
    
    # If text is short enough, process directly
    if len(words) <= CHUNK_SIZE:
        word_predictions, word_confidences = predict_chunk(words)
    else:
        # Process long text in overlapping chunks
        word_predictions = []
        word_confidences = []
        
        i = 0
        while i < len(words):
            # Get chunk with overlap
            chunk_end = min(i + CHUNK_SIZE, len(words))
            chunk_words = words[i:chunk_end]
            
            # Predict for chunk
            chunk_preds, chunk_confs = predict_chunk(chunk_words)
            
            # For overlapping regions, we only keep predictions from the first occurrence
            if i == 0:
                # First chunk: keep all predictions
                word_predictions.extend(chunk_preds)
                word_confidences.extend(chunk_confs)
                i += CHUNK_SIZE
            else:
                # Subsequent chunks: skip overlap region
                word_predictions.extend(chunk_preds[OVERLAP:])
                word_confidences.extend(chunk_confs[OVERLAP:])
                i += CHUNK_SIZE
            
            # If we've processed all words, break
            if chunk_end >= len(words):
                break
    
    # Construct output text
    output_words = []
    punctuation_details = []
    
    for i, (word, pred, conf) in enumerate(zip(words, word_predictions, word_confidences)):
        output_words.append(word)
        
        punct = id2punct[pred]
        if punct:
            output_words.append(punct)
            punctuation_details.append({
                'position': i + 1,
                'word': word,
                'punctuation': punct,
                'label': id2label[pred],
                'confidence': conf
            })
    
    # Create output text
    output_text = ' '.join(output_words)
    
    # Create detailed output if requested
    details_text = ""
    if show_details and punctuation_details:
        details_text = f"### 📊 Punctuation Details (Total words: {len(words)}):\n\n"
        details_text += "| Position | After Word | Punctuation | Type | Confidence |\n"
        details_text += "|----------|------------|-------------|------|------------|\n"
        for p in punctuation_details:
            details_text += f"| {p['position']} | {p['word']} | {p['punctuation']} | {p['label']} | {p['confidence']:.3f} |\n"
    elif show_details:
        details_text = f"ℹ️ No punctuation was added to this text. (Processed {len(words)} words)"
    
    return output_text, details_text


def batch_predict(text: str):
    """
    Process multiple sentences (one per line).
    """
    if not text or not text.strip():
        return "⚠️ Please enter some text!"
    
    lines = [line.strip() for line in text.split('\n') if line.strip()]
    results = []
    
    for i, line in enumerate(lines, 1):
        output, _ = predict_punctuation(line, show_details=False)
        results.append(f"{i}. {output}")
    
    return '\n\n'.join(results)


# ============================================================================
# GRADIO INTERFACE
# ============================================================================

# Example sentences with varying lengths
examples = [
    # Short examples (10-30 words)
    ["আমি স্কুলে যাই তুমি কোথায় যাও"],
    ["এটা কি তোমার বই"],
    ["দারুণ এটা তো অসাধারণ"],
    ["তুমি কেমন আছ আমি ভালো আছি"],
    # ~100 words - Daily life story
    ["আজ সকালে আমি ঘুম থেকে উঠে দেখলাম আকাশে মেঘ জমেছে বৃষ্টি হতে পারে আমি তাড়াতাড়ি তৈরি হয়ে নাশতা করলাম মা আমাকে ভাত ডাল আর মাছ দিয়েছিল খুব সুস্বাদু ছিল খাওয়ার পর আমি স্কুলের জন্য বের হলাম রাস্তায় অনেক যানজট ছিল অনেক দেরি হয়ে গেল ক্লাসে পৌঁছতে শিক্ষক একটু রাগ করলেন কিন্তু আমি কারণ বললে তিনি বুঝলেন আজকের ক্লাসে আমরা বাংলা সাহিত্য পড়লাম রবীন্দ্রনাথ ঠাকুরের কবিতা খুব ভালো লাগল"],
    ]

# Custom CSS
custom_css = """
#output_text {
    font-size: 18px;
    line-height: 1.8;
    padding: 15px;
    border-radius: 8px;
    background-color: #f8f9fa;
}
#input_text {
    font-size: 16px;
    line-height: 1.6;
}
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
"""

# Create Gradio interface with tabs
with gr.Blocks(css=custom_css, theme=gr.themes.Default(), title="Bangla Punctuation Restoration") as demo:
    
    gr.Markdown(
        """
        # 🇧🇩 Bangla Punctuation Restoration
        
        Add punctuation marks to unpunctuated Bangla text using AI. 
        This model can add: comma (,), dari (।), question mark (?), exclamation (!), semicolon (;), colon (:), and hyphen (-).
        
        **✨ Now supports texts of any length!** Long texts are automatically processed in chunks.
        
        **Model:** [Badhon/Bangla-Punc-Restore-Model](https://huggingface.co/Badhon/Bangla-Punc-Restore-Model)
        """
    )
    
    with gr.Tabs():
        # Tab 1: Single Text
        with gr.Tab("✍️ Single Text"):
            gr.Markdown("### Enter unpunctuated Bangla text (any length):")
            
            with gr.Row():
                with gr.Column(scale=1):
                    input_text = gr.Textbox(
                        label="Input Text (without punctuation)",
                        placeholder="আমি স্কুলে যাই তুমি কোথায় যাও",
                        lines=8,
                        elem_id="input_text"
                    )
                    
                    show_details = gr.Checkbox(
                        label="Show detailed punctuation information",
                        value=False
                    )
                    
                    predict_btn = gr.Button("🔄 Restore Punctuation", variant="primary", size="lg")
            
            output_text = gr.Textbox(
                label="Output Text (with punctuation)",
                lines=8,
                elem_id="output_text"
            )
            
            details_output = gr.Markdown(label="Details")
            
            predict_btn.click(
                fn=predict_punctuation,
                inputs=[input_text, show_details],
                outputs=[output_text, details_output]
            )
            
            gr.Markdown("### 📝 Examples:")
            gr.Examples(
                examples=examples,
                inputs=input_text,
                label="Click an example to try:"
            )
        
        # Tab 2: Batch Processing
        with gr.Tab("📚 Batch Processing"):
            gr.Markdown(
                """
                ### Process multiple sentences at once
                Enter one sentence per line. Each sentence will be processed separately.
                """
            )
            
            batch_input = gr.Textbox(
                label="Input Sentences (one per line)",
                placeholder="আমি স্কুলে যাই\nতুমি কোথায় যাও\nএটা কি তোমার বই",
                lines=10
            )
            
            batch_btn = gr.Button("🔄 Process All", variant="primary", size="lg")
            
            batch_output = gr.Textbox(
                label="Processed Sentences",
                lines=10,
                elem_id="output_text"
            )
            
            batch_btn.click(
                fn=batch_predict,
                inputs=batch_input,
                outputs=batch_output
            )
        
        # Tab 3: About
        with gr.Tab("ℹ️ About"):
            gr.Markdown(
                f"""
                ## About This App
                
                This application uses a fine-tuned transformer model to automatically add punctuation marks to Bangla text.
                
                ### ✨ New Features:
                - **Handles long texts**: Automatically processes texts longer than {CHUNK_SIZE} words using smart chunking
                - **Context preservation**: Uses overlapping windows to maintain context across chunks
                - **No length limit**: Process texts of any length!
                
                ### Supported Punctuation Marks:
                - **Comma (,)**: Used to separate clauses or items in a list
                - **Dari (।)**: Bangla full stop, marks the end of a sentence
                - **Question Mark (?)**: Indicates a question
                - **Exclamation Mark (!)**: Expresses strong emotion or emphasis
                - **Semicolon (;)**: Connects closely related independent clauses
                - **Colon (:)**: Introduces lists, quotes, or explanations
                - **Hyphen (-)**: Connects words or indicates breaks
                
                ### How It Works:
                1. Long texts are split into chunks of ~{CHUNK_SIZE} words with {OVERLAP}-word overlap
                2. Each chunk is analyzed word by word
                3. The model predicts whether punctuation should follow each word
                4. Results are combined seamlessly
                
                ### Model Details:
                - **Model**: Token Classification (BERT-based)
                - **Framework**: Hugging Face Transformers
                - **Language**: Bangla (Bengali)
                - **Task**: Punctuation Restoration
                - **Max Model Length**: {MAX_LENGTH} tokens
                
                ### Tips for Best Results:
                - Use properly spelled Bangla words
                - Avoid mixing languages in a single sentence
                - Remove any existing punctuation from input text
                - The model works well with texts of any length!
                
                ### Credits:
                Developed by Badhon | Model available on [Hugging Face](https://huggingface.co/Badhon/Bangla-Punc-Restore-Model)
                """
            )
    
    gr.Markdown(
        """
        ---
        💡 **Note**: This is an AI model and may not always be 100% accurate. Please review the output for important documents.
        """
    )

# ============================================================================
# LAUNCH APP
# ============================================================================
if __name__ == "__main__":
    demo.launch(
        share=True,
        show_error=True
    )