Update Text_Summary_app.py

Text_Summary_app.py

@@ -1,30 +1,46 @@
-import gradio as gr
 import os
+import io
+from IPython.display import Image, display, HTML
+from PIL import Image
+import base64
 from transformers import pipeline, AutoTokenizer
 
 # Initialize summarizer and tokenizer
-summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", tokenizer="sshleifer/distilbart-cnn-12-6")
+summarizer = pipeline ("summarization", model="sshleifer/distilbart-cnn-12-6", tokenizer="sshleifer/distilbart-cnn-12-6")
 tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
 
+import json
+
 def summarize_text(input_text):
-    """Summarizes the given input text."""
+    """Summarizes the given input text.
+
+    Args:
+        input_text (str): The text to be summarized.
+
+    Returns:
+        dict: A dictionary containing the summary under the 'summary' key.
+    """
+
+    # Tokenize and truncate input if necessary
     max_length = tokenizer.model_max_length
     inputs = tokenizer(input_text, truncation=True, max_length=max_length, return_tensors="pt")
+
+    # Generate summary
     summary_ids = summarizer.model.generate(inputs.input_ids, max_length=50, min_length=10, do_sample=False)
     summary_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+
+    # Return summary as a dictionary
     return {"summary": summary_text}
 
 def generate_summary(input):
     output = summarize_text(input)
-    return output["summary"]  # Return the summary directly
-
+    return output
+    
 gr.close_all()
-demo = gr.Interface(
-    fn=generate_summary,
-    inputs=[gr.Textbox(label="Text to summarize", lines=6)],
-    outputs=[gr.Textbox(label="Summary", lines=3)],
-    title="Text Summarization",
-    description="Summarize text using the 'shleifer/distilbart-cnn-12-6' language model.",
-)
-
-demo.launch()
+demo = gr.Interface(fn=generate_summary, 
+                        inputs=[gr.Textbox(label="Text to summarize", lines=6)], 
+                        outputs=[gr.Textbox(label="Summary", lines=3)],
+                        title="Text Summarization",
+                        description="Summarize text using the 'shleifer/distilbart-cnn-12-6' language model.",                    
+                    )
+demo.launch(share=True, server_port=int(os.environ['PORT']))