app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import os

# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Load model and tokenizer
def load_model():
    """
    Load the T5 model from safetensors format
    """
    try:
        print("Loading model from safetensors format...")
        
        # Load tokenizer and model from current directory (root of Space)
        tokenizer = AutoTokenizer.from_pretrained(".")
        model = AutoModelForSeq2SeqLM.from_pretrained(".")
        
        print("✓ Model and tokenizer loaded successfully")
        print(f"  Model parameters: {model.num_parameters():,}")
        print(f"  Device: {device}")
        
        # Create summarization pipeline
        summarizer = pipeline(
            "summarization",
            model=model,
            tokenizer=tokenizer,
            device=0 if device == "cuda" else -1
        )
        
        return summarizer
    
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        print("Make sure the model files (model.safetensors, config.json, tokenizer.json) are in the root directory")
        print(f"Current directory contents: {os.listdir('.')}")
        return None

# Load model at startup
print("Initializing Text Summarization Model...")
summarizer = load_model()

def summarize_text(
    text,
    max_length=128,
    min_length=30,
    num_beams=4,
    length_penalty=1.0,
    temperature=1.0,
    top_p=0.9,
    do_sample=False,
    repetition_penalty=1.0
):
    """
    Summarize the input text with customizable parameters
    """
    if summarizer is None:
        return "❌ Error: Model not loaded. Please check model files (model.safetensors, config.json, tokenizer.json) in the root directory."
    
    if not text.strip():
        return "⚠️ Please enter some text to summarize."
    
    try:
        # Generate summary with parameters
        summary = summarizer(
            text,
            max_length=int(max_length),
            min_length=int(min_length),
            num_beams=int(num_beams),
            length_penalty=float(length_penalty),
            temperature=float(temperature) if do_sample else 1.0,
            top_p=float(top_p) if do_sample else 1.0,
            do_sample=do_sample,
            repetition_penalty=float(repetition_penalty),
            early_stopping=True,
            truncation=True
        )
        
        return summary[0]['summary_text']
    
    except Exception as e:
        return f"❌ Error generating summary: {str(e)}"

def analyze_text(text):
    """
    Analyze the input text and provide statistics
    """
    if not text.strip():
        return "No text provided"
    
    words = len(text.split())
    sentences = text.count('.') + text.count('!') + text.count('?')
    chars = len(text)
    
    analysis = f"""
📊 **Text Analysis:**
• **Words:** {words:,}
• **Sentences:** {sentences}
• **Characters:** {chars:,}
• **Avg words/sentence:** {words/sentences:.1f} (if sentences > 0 else 0)
• **Est. reading time:** {words/200:.1f} minutes
    """
    
    return analysis.strip()

# Create Gradio interface
with gr.Blocks(
    title="T5 Text Summarization",
    theme=gr.themes.Soft(),
    css="""
        .output-text {font-size: 16px !important;}
        .input-text {font-size: 14px !important;}
    """
) as demo:
    
    gr.Markdown("""
    # 🤖 Advanced Text Summarization with T5
    ### Fine-tuned on CNN/DailyMail Dataset
    
    This app uses a fine-tuned **T5-small** model for abstractive text summarization.
    Upload your text and get concise, coherent summaries with advanced controls.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📝 Input Text")
            text_input = gr.Textbox(
                label="Enter text to summarize",
                placeholder="Paste your article or text here...",
                lines=12,
                max_lines=20
            )
            
            analyze_btn = gr.Button("📊 Analyze Text", variant="secondary", size="sm")
            analysis_output = gr.Markdown(
                label="Text Analysis",
                value=""
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Summarization Settings")
            
            with gr.Accordion("📏 Length Control", open=True):
                max_length = gr.Slider(
                    minimum=30,
                    maximum=300,
                    value=128,
                    step=10,
                    label="Max Summary Length (tokens)",
                    info="Maximum length of generated summary"
                )
                min_length = gr.Slider(
                    minimum=10,
                    maximum=100,
                    value=30,
                    step=5,
                    label="Min Summary Length (tokens)",
                    info="Minimum length of generated summary"
                )
            
            with gr.Accordion("🎯 Generation Parameters", open=False):
                num_beams = gr.Slider(
                    minimum=1,
                    maximum=8,
                    value=4,
                    step=1,
                    label="Num Beams",
                    info="Beam search width (higher = better quality, slower)"
                )
                
                length_penalty = gr.Slider(
                    minimum=0.5,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="Length Penalty",
                    info="Encourage longer (>1.0) or shorter (<1.0) summaries"
                )
                
                repetition_penalty = gr.Slider(
                    minimum=1.0,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="Repetition Penalty",
                    info="Reduce repetition (1.0 = no penalty)"
                )
            
            with gr.Accordion("🎲 Sampling (Advanced)", open=False):
                do_sample = gr.Checkbox(
                    value=False,
                    label="Enable Sampling",
                    info="Use stochastic sampling instead of greedy/beam search"
                )
                
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="Temperature",
                    info="Creativity (higher = more creative, only with sampling)"
                )
                
                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.9,
                    step=0.05,
                    label="Top-p (Nucleus Sampling)",
                    info="Probability mass cutoff (only with sampling)"
                )
            
            summarize_btn = gr.Button("🚀 Generate Summary", variant="primary", size="lg")
    
    with gr.Row():
        summary_output = gr.Textbox(
            label="📋 Generated Summary",
            placeholder="Your summary will appear here...",
            lines=8,
            max_lines=15,
            show_copy_button=True,
            elem_classes="output-text"
        )
    
    # Event handlers
    analyze_btn.click(
        fn=analyze_text,
        inputs=text_input,
        outputs=analysis_output
    )
    
    summarize_btn.click(
        fn=summarize_text,
        inputs=[
            text_input,
            max_length,
            min_length,
            num_beams,
            length_penalty,
            temperature,
            top_p,
            do_sample,
            repetition_penalty
        ],
        outputs=summary_output
    )
    
    # Examples
    gr.Examples(
        examples=[
            ["Artificial intelligence (AI) has revolutionized numerous industries in recent years. From healthcare to finance, AI systems are now capable of performing tasks that once required human intelligence. Machine learning algorithms can analyze vast amounts of data to identify patterns and make predictions. Natural language processing enables computers to understand and generate human language. Computer vision allows machines to interpret and understand visual information from the world. Despite these advances, concerns about AI ethics, bias, and job displacement remain significant challenges that society must address as the technology continues to evolve."],
            
            ["The COVID-19 pandemic has had a profound impact on global health and economies. Countries worldwide implemented lockdowns and social distancing measures to curb the spread of the virus. Healthcare systems were overwhelmed, and millions of lives were lost. Economic activities came to a standstill, leading to widespread unemployment and financial hardship. Vaccination campaigns were launched globally, providing hope for controlling the pandemic. However, new variants continue to emerge, requiring ongoing vigilance and adaptation of public health strategies. The pandemic has fundamentally changed how people work, learn, and interact with each other."],
            
            ["Climate change represents one of the most significant challenges facing humanity. Rising global temperatures are causing extreme weather events, sea level rise, and ecosystem disruption. The burning of fossil fuels, deforestation, and industrial activities are major contributors to greenhouse gas emissions. International agreements like the Paris Accord aim to limit global warming to well below 2 degrees Celsius. Renewable energy sources such as solar and wind power offer sustainable alternatives. Individual actions, corporate responsibility, and government policies are all crucial in addressing this global crisis. Scientists warn that immediate action is needed to prevent irreversible damage to our planet."]
        ],
        inputs=text_input,
        label="📚 Example Articles"
    )
    
    gr.Markdown("""
    ---
    ### 🎯 Parameter Guide:
    
    - **Max/Min Length**: Control the summary length in tokens (~0.75 words per token)
    - **Num Beams**: Beam search improves quality but is slower (4 is recommended)
    - **Length Penalty**: >1.0 encourages longer summaries, <1.0 prefers shorter ones
    - **Repetition Penalty**: Reduces repetition (increase if summary repeats itself)
    - **Temperature**: Controls randomness when sampling is enabled (higher = more creative)
    - **Top-p**: Nucleus sampling threshold (only with sampling enabled)
    
    ### 📊 Model Information:
    
    - **Base Model**: T5-small (60M parameters)
    - **Fine-tuned on**: CNN/DailyMail dataset
    - **Task**: Abstractive text summarization
    - **Format**: Safetensors (efficient & secure)
    - **Evaluation**: ROUGE-1, ROUGE-2, ROUGE-L metrics
    
    ### 💡 Tips for Best Results:
    
    1. **Optimal length**: 200-1000 words for input text
    2. **Clear structure**: Well-formatted articles work best
    3. **Adjust length**: Set max_length based on desired summary length
    4. **Use beam search**: Keep num_beams=4 for best quality
    5. **Avoid sampling**: For news articles, disable sampling for factual summaries
    """)

if __name__ == "__main__":
    # Launch the app
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )