Update app.py

app.py

@@ -1,45 +1,29 @@
 import streamlit as st
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from pydrive.auth import GoogleAuth
-from pydrive.drive import GoogleDrive
-import os
-from io import BytesIO
-from zipfile import ZipFile
-
-# Initialize Google Auth and Drive
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, GenerationConfig
+from pydrive2.auth import GoogleAuth
+from pydrive2.drive import GoogleDrive
+
+# Authenticate and create the PyDrive client.
 gauth = GoogleAuth()
-gauth.LocalWebserverAuth()  # Authenticates and opens a browser window
+gauth.LocalWebserverAuth()  # Creates a local webserver and automatically handles authentication.
 drive = GoogleDrive(gauth)
 
-# Streamlit UI
-st.title("Text Summarizer")
-
-# Enter the file ID of your model.zip on Google Drive
-model_file_id = st.text_input("Enter the Google Drive file ID of the model.zip")
+# Update this path to your local path where the model is stored
+model_path = '/content/drive/My Drive/bart-base'
 
-if model_file_id:
-    try:
-        # Download the file from Google Drive
-        downloaded = drive.CreateFile({'id': model_file_id}).GetContentString()
-        
-        # Load the model from the downloaded zip file
-        with ZipFile(BytesIO(downloaded.encode()), 'r') as zip_ref:
-            zip_ref.extractall("model_directory")
-
-        # Load the model from the extracted directory
-        model_path = "model_directory"
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
-        st.success("Model loaded successfully!")
-    except Exception as e:
-        st.error(f"Failed to load model: {e}")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+except Exception as e:
+    st.error(f"Failed to load model: {e}")
 
-# Text area for input
+# Streamlit UI
+st.title("Text Summarizer")
 text = st.text_area("Enter the text to generate its Summary:")
 
 # Configuration for generation
-generation_config = {'max_length': 100, 'do_sample': True, 'temperature': 0.7}
+generation_config = GenerationConfig(max_new_tokens=100, do_sample=True, temperature=0.7)
 
 if text:
     try:
@@ -48,14 +32,14 @@ if text:
 
         # Generate output
         with torch.no_grad():
-            model_output = model.generate(inputs_encoded["input_ids"], **generation_config)[0]
+            model_output = model.generate(inputs_encoded["input_ids"], generation_config=generation_config)[0]
 
         # Decode output
         output = tokenizer.decode(model_output, skip_special_tokens=True)
 
-        # Display results
+        # Display results in a box with a title
         with st.expander("Output", expanded=True):
             st.write(output)
 
     except Exception as e:
-        st.error(f"An error occurred during summarization: {e}")
+        st.error(f"An error occurred during summarization: {e}")