Update app.py

app.py

@@ -3,9 +3,31 @@ import torch
 from transformers import AutoTokenizer, AlbertForSequenceClassification
 import numpy as np
 import os
+import gdown
 
-# Define paths to saved models
+# Define Google Drive folder IDs for each model
+model_drive_ids = {
+    "sentiment": "your_sentiment_folder_id",  # Replace with actual folder ID
+    "emotion": "your_emotion_folder_id",      # Replace with actual folder ID
+    "hate_speech": "your_hate_speech_folder_id",  # Replace with actual folder ID
+    "sarcasm": "your_sarcasm_folder_id"       # Replace with actual folder ID
+}
+
+# Define local directory to store downloaded models
 save_dir = "./saved_models"
+os.makedirs(save_dir, exist_ok=True)
+
+# Download models from Google Drive
+for task, folder_id in model_drive_ids.items():
+    output_dir = os.path.join(save_dir, task)
+    if not os.path.exists(output_dir):
+        gdown.download_folder(
+            f"https://drive.google.com/drive/folders/1kEXKoJxxD5-0FO8WvtagzseSIC5q-rRY?usp=sharing/{folder_id}",
+            output=output_dir,
+            quiet=False
+        )
+
+# Define model paths
 tasks = ["sentiment", "emotion", "hate_speech", "sarcasm"]
 model_paths = {task: f"{save_dir}/{task}" for task in tasks}
 
@@ -24,7 +46,7 @@ tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indic-bert")
 models = {}
 for task in tasks:
     if not os.path.exists(model_paths[task]):
-        raise FileNotFoundError(f"Model directory {model_paths[task]} not found. Ensure saved_models/{task} exists.")
+        raise FileNotFoundError(f"Model directory {model_paths[task]} not found.")
     models[task] = AlbertForSequenceClassification.from_pretrained(model_paths[task])
 
 # Function to predict for a single task