Update app.py

app.py

@@ -24,21 +24,12 @@ def load_model(model_name):
     global summarizer, tokenizer, max_tokens
     try:
         # Load the summarization pipeline with the selected model
-        summarizer = pipeline("summarization", model=model_name, torch_dtype=torch.bfloat16)
-        # Load the tokenizer for the selected model
+        summarizer = pipeline("summarization", model=model_name, torch_dtype=torch.float32)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        # Load the configuration for the selected model
         config = AutoConfig.from_pretrained(model_name)
 
-        # Determine the maximum tokens based on available configuration attributes
-        if hasattr(config, 'max_position_embeddings'):
-            max_tokens = config.max_position_embeddings
-        elif hasattr(config, 'n_positions'):
-            max_tokens = config.n_positions
-        elif hasattr(config, 'd_model'):
-            max_tokens = config.d_model  # for T5 models, d_model is a rough proxy
-        else:
-            max_tokens = "Unknown"
+        # Set a reasonable default for max_tokens if not available
+        max_tokens = getattr(config, 'max_position_embeddings', 1024)
 
         return f"Model {model_name} loaded successfully! Max tokens: {max_tokens}"
     except Exception as e:
@@ -50,49 +41,41 @@ def summarize_text(input, min_length, max_length):
     if summarizer is None:
         return "No model loaded!"
 
-    # Tokenize the input text and check the number of tokens
-    input_tokens = tokenizer.encode(input, return_tensors="pt")
-    num_tokens = input_tokens.shape[1]
-    if num_tokens > max_tokens:
-        # Return an error message if the input text exceeds the maximum token limit
-        return f"Error: The input text has {num_tokens} tokens, which exceeds the maximum allowed {max_tokens} tokens. Please enter shorter text."
-
-    # Calculate minimum and maximum summary length based on the percentages
-    min_summary_length = int(num_tokens * (min_length / 100))
-    max_summary_length = int(num_tokens * (max_length / 100))
-
-    # Summarize the input text using the loaded model with specified lengths
-    output = summarizer(input, min_length=min_summary_length, max_length=max_summary_length)
-    return output[0]['summary_text']
+    try:
+        # Tokenize the input text and check the number of tokens
+        input_tokens = tokenizer.encode(input, return_tensors="pt")
+        num_tokens = input_tokens.shape[1]
+        if num_tokens > max_tokens:
+            return f"Error: Input exceeds the max token limit of {max_tokens}."
+
+        # Ensure min/max lengths are within bounds
+        min_summary_length = max(10, int(num_tokens * (min_length / 100)))
+        max_summary_length = min(max_tokens, int(num_tokens * (max_length / 100)))
+
+        # Summarize the input text
+        output = summarizer(input, min_length=min_summary_length, max_length=max_summary_length, truncation=True)
+        return output[0]['summary_text']
+    except Exception as e:
+        return f"Summarization failed: {str(e)}"
 
 
 # Gradio Interface
 with gr.Blocks() as demo:
     with gr.Row():
-        # Dropdown menu for selecting the model
         model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="sshleifer/distilbart-cnn-12-6")
-        # Button to load the selected model
         load_button = gr.Button("Load Model")
 
-    # Textbox to display the load status
     load_message = gr.Textbox(label="Load Status", interactive=False)
 
-    # Slider for minimum summary length
     min_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Minimum Summary Length (%)", value=10)
-    # Slider for maximum summary length
     max_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Maximum Summary Length (%)", value=20)
 
-    # Textbox for inputting the text to be summarized
     input_text = gr.Textbox(label="Input text to summarize", lines=6)
-    # Button to trigger the summarization
     summarize_button = gr.Button("Summarize Text")
-    # Textbox to display the summarized text
     output_text = gr.Textbox(label="Summarized text", lines=4)
 
-    # Define the actions for the load button and summarize button
     load_button.click(fn=load_model, inputs=model_dropdown, outputs=load_message)
     summarize_button.click(fn=summarize_text, inputs=[input_text, min_length_slider, max_length_slider],
                            outputs=output_text)
 
-# Launch the Gradio interface
 demo.launch()