Update app.py

app.py

@@ -1,36 +1,34 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import PegasusTokenizer, PegasusForConditionalGeneration
 import gradio as gr
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("mynkchaudhry/Summarization-Pro",force_download=True,cache_dir=None)
-model = AutoModelForSeq2SeqLM.from_pretrained("mynkchaudhry/Summarization-Pro",force_download=True,cache_dir=None)
+# Load model and tokenizer based on the provided configuration
+tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-cnn_dailymail")
+model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-cnn_dailymail")
 
-def summarize(text):
-    inputs = tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
-    summary_ids = model.generate(
-        inputs.input_ids,
-        max_length=1500,  # Increased max_length for longer summaries
-        min_length=100,   # Ensures summaries are sufficiently long
-        length_penalty=0.9,  # Encourages longer summaries but balances precision
-        num_beams=6,      # More beams for better exploration of possible summaries
-        early_stopping=True
-    )
+def summarize(text, prompt):
+    # Prepend the prompt to the input text
+    inputs = tokenizer(prompt + " " + text, return_tensors="pt", max_length=1024, truncation=True)
+    # Generate the summary
+    summary_ids = model.generate(inputs.input_ids, max_length=128, min_length=32, length_penalty=0.8, num_beams=8, early_stopping=True)
     summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return summary
 
 # Gradio interface
 interface = gr.Interface(
     fn=summarize,
-    inputs=gr.Textbox(lines=10, label="Input Text"),
+    inputs=[
+        gr.Textbox(lines=10, label="Input Text"),
+        gr.Textbox(lines=2, label="Prompt", placeholder="Enter a prompt for the summary here...")
+    ],
     outputs=gr.Textbox(label="Summary"),
     examples=[
-        ["The quick brown fox jumps over the lazy dog. This sentence is an English-language pangram—a sentence that contains all the letters of the English alphabet. It is often used for testing typewriters, keyboards, and computer fonts."],
-        ["Artificial intelligence (AI) is intelligence demonstrated by machines, in contrast to the natural intelligence displayed by humans and animals. Leading AI textbooks define the field as the study of intelligent agents: any device that perceives its environment and takes actions that maximize its chance of successfully achieving its goals."],
-        ["In data science, an anomaly is an observation that deviates significantly from other observations. Anomalies are also known as outliers, novelties, noise, deviations, and exceptions. Anomaly detection finds application in many domains, including fraud detection, network security, and fault detection."]
+        ["The quick brown fox jumps over the lazy dog. This sentence is an English-language pangram—a sentence that contains all the letters of the English alphabet. It is often used for testing typewriters, keyboards, and computer fonts.", "Summarize this text:"],
+        ["Artificial intelligence (AI) is intelligence demonstrated by machines, in contrast to the natural intelligence displayed by humans and animals. Leading AI textbooks define the field as the study of intelligent agents: any device that perceives its environment and takes actions that maximize its chance of successfully achieving its goals.", "Summarize this text:"],
+        ["In data science, an anomaly is an observation that deviates significantly from other observations. Anomalies are also known as outliers, novelties, noise, deviations, and exceptions. Anomaly detection finds application in many domains, including fraud detection, network security, and fault detection.", "Summarize this text:"]
     ],
     title="Text Summarization",
-    description="Generate a summary from the input text using a pre-trained summarization model.",
-    theme="freddyaboulton/dracula_revamped"  # Ensure this theme exists or use a default theme
+    description="Generate a summary from the input text using a pre-trained PEGASUS model. Optionally, provide a prompt to guide the summarization.",
+    theme="freddyaboulton/dracula_revamped"
 )
 
 interface.launch()