Update app.py

app.py

@@ -1,22 +1,34 @@
 import gradio as gr
-from transformers import BartTokenizer, BartForConditionalGeneration
+from transformers import LEDTokenizer, LEDForConditionalGeneration
 
-model_name = "facebook/bart-large-cnn"
-tokenizer = BartTokenizer.from_pretrained(model_name)
-model = BartForConditionalGeneration.from_pretrained(model_name)
+# Use Longformer Encoder-Decoder (LED) model
+model_name = "allenai/led-large-16384"
+tokenizer = LEDTokenizer.from_pretrained(model_name)
+model = LEDForConditionalGeneration.from_pretrained(model_name)
 
 def summarize_text(text):
-    inputs = tokenizer([text], max_length=1024, return_tensors="pt", truncation=True)
-    summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=128, early_stopping=True)
+    # Tokenize input with truncation to fit within 16,384 tokens
+    inputs = tokenizer([text], max_length=16384, return_tensors="pt", truncation=True)
+
+    # Generate summary with adjusted parameters
+    summary_ids = model.generate(
+        inputs["input_ids"],
+        num_beams=4,
+        max_length=512,  # Can be adjusted based on summary size needs
+        min_length=100,
+        early_stopping=True
+    )
+
     summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return summary
 
+# Gradio Interface
 iface = gr.Interface(
     fn=summarize_text,
     inputs="text",
     outputs="text",
-    title="BART Summarizer",
-    description="Enter text to get a summary using Facebook BART."
+    title="Longformer Summarizer",
+    description="Enter text to get a summary using the Longformer Encoder-Decoder."
 )
 
 if __name__ == "__main__":