Update app.py

app.py

@@ -1,18 +1,21 @@
 import os
-os.system('pip install streamlit transformers torch')
-
 import streamlit as st
-from transformers import pipeline
-import torch
+from transformers import PegasusTokenizer, PegasusForConditionalGeneration
 
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+# Install necessary packages (only needed if not already installed)
+os.system('pip install streamlit transformers torch')
 
 # Load the model and tokenizer
 model_path = '.'  # Path to the current directory where files are located
 
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
-summarizer = pipeline('summarization', model=model, tokenizer=tokenizer)
+tokenizer = PegasusTokenizer.from_pretrained(model_path)
+model = PegasusForConditionalGeneration.from_pretrained(model_path)
+
+def summarize_text(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="longest")
+    summary_ids = model.generate(inputs["input_ids"], max_length=150, min_length=30, do_sample=False)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
 
 st.title("Text Summarization with Fine-Tuned Model")
 st.write("Enter text to generate a summary using the fine-tuned summarization model.")
@@ -21,9 +24,9 @@ text = st.text_area("Input Text", height=200)
 if st.button("Summarize"):
     if text:
         with st.spinner("Summarizing..."):
-            summary = summarizer(text, max_length=150, min_length=30, do_sample=False)
+            summary = summarize_text(text)
             st.success("Summary Generated")
-            st.write(summary[0]['summary_text'])
+            st.write(summary)
     else:
         st.warning("Please enter some text to summarize.")
 
@@ -39,4 +42,4 @@ if __name__ == "__main__":
         </style>
         """,
         unsafe_allow_html=True
-    )
+    )