Update app.py

app.py

@@ -1,14 +1,14 @@
 import torch
 import gradio as gr
-from transformers import BartTokenizer, BartForConditionalGeneration
+from transformers import PegasusTokenizer, PegasusForConditionalGeneration
 
-# Define the BART model and tokenizer
-MODEL_NAME = 'VishnuPottabatthini/BART_demo'  # Change this to the model you want to use
+# Define the PEGASUS model and tokenizer
+MODEL_NAME = 'VishnuPottabatthini/PEGASUS_Large'  # Change this to the PEGASUS model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the tokenizer and model
-tokenizer = BartTokenizer.from_pretrained(MODEL_NAME)
-model = BartForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
+tokenizer = PegasusTokenizer.from_pretrained(MODEL_NAME)
+model = PegasusForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
 
 # Define the summarization function
 def summarize(text, state):
@@ -53,7 +53,7 @@ mf_summarize = gr.Interface(
     title="Article Summarization",
     live=True,
     description=(
-        "Enter a long piece of text to generate a concise summary using a BART model. "
+        "Enter a long piece of text to generate a concise summary using a PEGASUS model. "
         "This demo uses a custom PEGASUS model from 🤗 Transformers."
     )
 )