app.py

### STREAMLIT APP ####
import streamlit as st
from transformers import BertTokenizer, EncoderDecoderModel, EncoderDecoderConfig

model_ckpt = 'ardavey/bert2gpt-indosum'
tokenizer = BertTokenizer.from_pretrained(model_ckpt)
tokenizer.bos_token = tokenizer.cls_token
tokenizer.eos_token = tokenizer.sep_token

config = EncoderDecoderConfig.from_pretrained(model_ckpt)
config.early_stopping = True

try:
 model = EncoderDecoderModel.from_pretrained(model_ckpt, config=config)
except Exception as e:
 st.error(f"An error occurred while loading the model: {e}")

text = st.text_area('Enter an article to summarize:')

if text and len(text) > 0:
 # Ensure input is a string
 input_ids = tokenizer.encode(text, return_tensors='pt', padding='longest', truncation=True, max_length=512)
 
 summary_ids = model.generate(input_ids,
 min_length=40,
 max_length=200,
 num_beams=10,
 repetition_penalty=2.0,
 length_penalty=1.0,
 no_repeat_ngram_size=3,
 use_cache=True,
 do_sample=False,
 top_k=50,
 )

 summary_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
 # Check if summary is empty
 if not summary_text.strip():
 st.warning("The model couldn't generate a summary.")
 else:
 # Capitalize the first letter of the summary and after each period
 def capitalize_sentences(text):
 sentences = text.split('. ')
 capitalized_sentences = [sentence[0].upper() + sentence[1:] if sentence else sentence for sentence in sentences]
 return '. '.join(capitalized_sentences)
 
 # Correct any wrong terms using the replacement_dict
 replacement_dict = {
 "optiglain": "OptiGuard",
 "telkom university": "Telkom University",
 "menyerbut": "menyebut"
 }
 
 for wrong_term, correct_term in replacement_dict.items():
 summary_text = summary_text.replace(wrong_term, correct_term)
 
 summary_text = capitalize_sentences(summary_text)

 st.subheader("Generated Summary:")
 st.write(summary_text)