Update app.py

app.py

@@ -3,12 +3,16 @@ import torch
 import os
 import chardet
 import time
-from core.summarizer import NarrativeSummarizer # <-- Import the core logic
+from core.summarizer import NarrativeSummarizer
 
 # Models available
 MODEL_OPTIONS = [
     "facebook/bart-large-cnn",
+    "sshleifer/distilbart-cnn-12-6",
+    "google/pegasus-cnn_dailymail",
     "google/long-t5-local-base",
+    "t5-small",
+    "t5-base",
     "mistralai/Mistral-7B-v0.1"
 ]
 
@@ -20,7 +24,7 @@ PROMPT_OPTIONS = [
     "Custom Prompt"
 ]
 
-def run_app(file_obj, text_input, model_name, local_model_path, prompt_type, custom_prompt_text, iterations, batch_size, target_word_count, progress=gr.Progress()):
+def run_app(file_obj, text_input, model_name, local_model_path, prompt_type, custom_prompt_text, iterations, batch_size, progress=gr.Progress()):
     start_time = time.time()
     
     # Check if custom prompt is selected but not provided
@@ -73,8 +77,7 @@ def run_app(file_obj, text_input, model_name, local_model_path, prompt_type, cus
                 batch_summaries = summarizer.summarize_batch(
                     batch, 
                     prompt_type, 
-                    custom_prompt_text if prompt_type == "Custom Prompt" else None,
-                    target_word_count
+                    custom_prompt_text if prompt_type == "Custom Prompt" else None
                 )
                 batch = batch_summaries
             condensed_chunks.extend(batch)
@@ -91,8 +94,7 @@ def run_app(file_obj, text_input, model_name, local_model_path, prompt_type, cus
             final_summary = summarizer.summarize_batch(
                 [combined], 
                 prompt_type, 
-                custom_prompt_text if prompt_type == "Custom Prompt" else None,
-                target_word_count
+                custom_prompt_text if prompt_type == "Custom Prompt" else None
             )[0]
 
         end_time = time.time()
@@ -118,7 +120,10 @@ model_tip = gr.Markdown(
     """
     **Model Selection Tips:**
     - **facebook/bart-large-cnn:** Fast, general-purpose summarization for short to medium texts.
+    - **sshleifer/distilbart-cnn-12-6:** A smaller, faster version of BART.
+    - **google/pegasus-cnn_dailymail:** Excellent for abstractive summarization.
     - **google/long-t5-local-base:** Designed for long documents; better context handling.
+    - **t5-small/t5-base:** Efficient and versatile for various tasks.
     - **mistralai/Mistral-7B-v0.1:** High-quality nuanced summaries; resource-intensive.
     - **Custom Local Model:** Specify a path to a downloaded model (e.g., `./models/my-bart-model`).
     """
@@ -151,9 +156,6 @@ with gr.Blocks(css="#status-log { overflow-y: scroll; max-height: 200px; }") as
         with gr.Row():
             iterations_slider = gr.Slider(minimum=1, maximum=5, step=1, label="Summarization Iterations", value=1)
             batch_size_slider = gr.Slider(minimum=1, maximum=8, step=1, label="Batch Size (for GPU)", value=4)
-        
-        with gr.Row():
-            target_word_count_slider = gr.Slider(minimum=20, maximum=200, step=10, label="Target Summary Word Count", value=50)
 
     summarize_button = gr.Button("Summarize")
     
@@ -171,7 +173,7 @@ with gr.Blocks(css="#status-log { overflow-y: scroll; max-height: 200px; }") as
         outputs=summarize_button
     ).then(
         fn=run_app,
-        inputs=[file_input, text_input, model_dropdown, local_model_path_input, prompt_dropdown, custom_prompt_input, iterations_slider, batch_size_slider, target_word_count_slider],
+        inputs=[file_input, text_input, model_dropdown, local_model_path_input, prompt_dropdown, custom_prompt_input, iterations_slider, batch_size_slider],
         outputs=[output_text, status_log, download_button, summarize_button]
     )