Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import BartTokenizer, BartForConditionalGeneration, pipeline
+from transformers import AutoTokenizer, AutoModelForConditionalGeneration, pipeline
 import zipfile
 import os
 import nltk
@@ -8,7 +8,7 @@ import nltk
 nltk.download('punkt')
 from nltk.tokenize import sent_tokenize
 
-# Define the path to the saved model zip file
+# Define the path to the saved model zip file (ensure there is no extra space)
 zip_model_path = 'bart_model-20240724T051306Z-001.zip'
 
 # Define the directory to extract the model
@@ -19,7 +19,7 @@ with zipfile.ZipFile(zip_model_path, 'r') as zip_ref:
     zip_ref.extractall(model_dir)
 
 # After unzipping, the model should be in a specific directory, check the directory structure
-model_path = os.path.join(model_dir, 'Bart')
+model_path = os.path.join(model_dir, 'Bart_model')
 
 # Print out the model_path for debugging
 print("Model Path:", model_path)
@@ -28,9 +28,9 @@ print("Model Path:", model_path)
 if not os.path.exists(model_path):
     st.error(f"Model directory {model_path} does not exist or is incorrect.")
 else:
-    # Load the tokenizer and model
-    tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
-    model = BartForConditionalGeneration.from_pretrained(model_path)
+    # Load the tokenizer and model using AutoTokenizer and AutoModelForConditionalGeneration
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForConditionalGeneration.from_pretrained(model_path)
 
     # Create a summarization pipeline
     summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)