ui/summarization_ui.py

import gradio as gr
from utils.ner_helpers import is_llm_model
from typing import Dict, List, Any
from tasks.summarization import text_summarization

def summarization_ui():
    """Summarization UI component"""
    
    # Define models
    SUMMARY_MODELS = [
        "gemini-2.0-flash"  # Only allow gemini-2.0-flash for now
        # "gpt-4",
        # "claude-2",
        # "facebook/bart-large-cnn",
        # "t5-small",
        # "qwen/Qwen2.5-3B-Instruct"
    ]
    DEFAULT_MODEL = "gemini-2.0-flash"
    
    def summarize(text, model, summary_length, custom_instructions):
        """Process text for summarization"""
        if not text.strip():
            return "No text provided"
            
        use_llm = is_llm_model(model)
        result = text_summarization(
            text=text, 
            model=model, 
            summary_length=summary_length,
            use_llm=use_llm
        )
        
        # Lưu ý: custom_instructions sẽ được sử dụng trong tương lai khi API hỗ trợ
            
        return result
    
    # UI Components
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Input Text",
                lines=8,
                placeholder="Enter text to summarize...",
                elem_id="summary-input-text"
            )

            summary_length = gr.Radio(
                ["Short", "Medium", "Long"], 
                value="Medium", 
                label="Summary Length",
                elem_id="summary-length-radio"
            )
            model = gr.Dropdown(
                SUMMARY_MODELS,
                value=DEFAULT_MODEL,
                label="Model",
                interactive=True,
                elem_id="summary-model-dropdown"
            )
            custom_instructions = gr.Textbox(
                label="Custom Instructions (optional)",
                lines=2,
                placeholder="Add any custom instructions for the model...",
                elem_id="summary-custom-instructions"
            )

            btn = gr.Button("Summarize", variant="primary", elem_id="summary-btn")
        
        with gr.Column():
            output = gr.Textbox(
                label="Summary", 
                lines=10,
                elem_id="summary-output"
            )
            
            # with gr.Accordion("About Summarization", open=False):
            #     gr.Markdown("""
            #     ## Text Summarization
                
            #     Text summarization condenses a document while preserving key information. This tool offers:
                
            #     - **Length control**: Choose between short, medium, or long summaries
            #     - **Multiple models**: Select from LLMs (like Gemini and GPT) or traditional models
            #     - **Custom instructions**: Tailor the summarization to your specific needs
                
            #     ### How it works
                
            #     - **LLM models** process your text using natural language understanding
            #     - **Traditional models** use extractive or abstractive techniques to identify and condense key information
                
            #     For best results with long texts, try different summary lengths to find the right balance between brevity and detail.
            #     """)
    
    # Event handlers
    btn.click(
        summarize, 
        inputs=[input_text, model, summary_length, custom_instructions], 
        outputs=output
    )
    
    return None