Update app.py

app.py

@@ -1,50 +1,77 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import gdown
 import os
 
 # Set the title of the Streamlit app
-st.title("Text Classification with Hugging Face Transformers")
+st.title("Text Summarization with Fine-Tuned BART")
 
 # Function to download the model from Google Drive
 def download_model_from_drive(file_id, dest_path):
     url = f'https://drive.google.com/uc?id={file_id}'
-    gdown.download(url, dest_path, quiet=False)
+    try:
+        gdown.download(url, dest_path, quiet=False)
+        st.success(f"Downloaded {dest_path}")
+    except Exception as e:
+        st.error(f"Error downloading {dest_path}: {e}")
+
+# Ensure the model directory exists
+model_dir = 'model'
+if not os.path.exists(model_dir):
+    os.makedirs(model_dir)
+
+# File IDs for your model components
+file_ids = {
+    'model': '1-V2bEtPR9Y3iBXK9zOR-qM5y9hKiQUnF',
+    'config': '1-T2etSP_k_3j5LzunWq8viKGQCQ5RMr_',
+    'tokenizer': '1-cRYNPWqlNNGRxeztympRRfVuy3hWuMY',
+    'vocab': '1-t9AhomeH7YIIpAqCGTok8wjvl0tml0F',
+    'merges': '1-l77_KEdK7GBFjMX_6UXGE-ZTGDraaDm'
+}
 
 # Download the model files
 with st.spinner("Downloading model..."):
-    download_model_from_drive('1-V2bEtPR9Y3iBXK9zOR-qM5y9hKiQUnF', 'model/model.safetensors')
-    download_model_from_drive('1-T2etSP_k_3j5LzunWq8viKGQCQ5RMr_', 'model/config.json')
-    download_model_from_drive('1-cRYNPWqlNNGRxeztympRRfVuy3hWuMY', 'model/tokenizer.json')
-    download_model_from_drive('1-t9AhomeH7YIIpAqCGTok8wjvl0tml0F', 'model/vocab.json')
-    download_model_from_drive('1-l77_KEdK7GBFjMX_6UXGE-ZTGDraaDm', 'model/merges.txt')
+    download_model_from_drive(file_ids['model'], os.path.join(model_dir, 'pytorch_model.bin'))
+    download_model_from_drive(file_ids['config'], os.path.join(model_dir, 'config.json'))
+    download_model_from_drive(file_ids['tokenizer'], os.path.join(model_dir, 'tokenizer.json'))
+    download_model_from_drive(file_ids['vocab'], os.path.join(model_dir, 'vocab.json'))
+    download_model_from_drive(file_ids['merges'], os.path.join(model_dir, 'merges.txt'))
 
 # Load the model and tokenizer
 @st.cache(allow_output_mutation=True)
 def load_model_and_tokenizer():
-    tokenizer = AutoTokenizer.from_pretrained('model')
-    # For Safetensors, you might need a custom loading mechanism
-    model = AutoModelForSequenceClassification.from_pretrained('model', use_safetensors=True)  # Adjust if necessary
-    return tokenizer, model
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+        return tokenizer, model
+    except Exception as e:
+        st.error(f"Error loading model or tokenizer: {e}")
+        return None, None
 
 tokenizer, model = load_model_and_tokenizer()
 
 # Input text from user
-input_text = st.text_area("Enter the text to classify:")
+input_text = st.text_area("Enter the text to summarize:")
 
-if st.button("Classify"):
+if st.button("Summarize"):
     if input_text:
-        # Tokenize the input text
-        inputs = tokenizer(input_text, return_tensors="pt")
-        
-        # Perform classification
-        with torch.no_grad():
-            outputs = model(**inputs)
-        
-        # Get the predicted class
-        predicted_class = torch.argmax(outputs.logits, dim=1).item()
-
-        st.write(f"Predicted Class: {predicted_class}")
+        if tokenizer and model:
+            try:
+                # Tokenize the input text
+                inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
+                
+                # Perform summarization
+                with torch.no_grad():
+                    summary_ids = model.generate(inputs['input_ids'], max_length=150, num_beams=4, early_stopping=True)
+                
+                # Decode the summary
+                summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+                
+                st.write(f"Summary: {summary}")
+            except Exception as e:
+                st.error(f"Error during summarization: {e}")
+        else:
+            st.error("Model or tokenizer not loaded.")
     else:
-        st.write("Please enter some text to classify.")
+        st.write("Please enter some text to summarize.")