Update app.py

app.py

@@ -1,36 +1,48 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import os
-
-# Set the path to your local model directory
-model_path = "./bart_samsum"
-
-# Check if the model path exists
-if not os.path.exists(model_path):
-    st.error(f"The path {model_path} does not exist. Please check the path.")
-else:
-    # Load the tokenizer and model from the local directory
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
-
-    # Streamlit app UI
-    st.title("BART Summarization Model")
-
-    input_text = st.text_area("Input Text", "Enter text here...")
-
-    if st.button("Generate Summary"):
-        if not input_text.strip():
-            st.warning("Please enter some text to summarize.")
-        else:
-            # Tokenize and generate summary
-            inputs = tokenizer(input_text, return_tensors="pt")
-            summary_ids = model.generate(inputs["input_ids"], max_length=150, num_beams=4, early_stopping=True)
-            summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-            
-            # Display the summary
-            st.subheader("Generated Summary")
-            st.write(summary)
-
-# Optionally, you can add a section to display model information or statistics
-st.sidebar.title("Model Information")
-st.sidebar.write("This app uses a fine-tuned BART model for summarization.")
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+from google_drive_downloader import GoogleDriveDownloader as gdd
+
+# Set the title of the Streamlit app
+st.title("Text Classification with Hugging Face Transformers")
+
+# Function to download the model from Google Drive
+def download_model_from_drive(file_id, dest_path):
+    gdd.download_file_from_google_drive(file_id=file_id, dest_path=dest_path, unzip=False)
+
+# 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')
+
+# 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
+
+tokenizer, model = load_model_and_tokenizer()
+
+# Input text from user
+input_text = st.text_area("Enter the text to classify:")
+
+if st.button("Classify"):
+    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}")
+    else:
+        st.write("Please enter some text to classify.")