Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AlbertForSequenceClassification
+from transformers import AutoTokenizer, AlbertForSequenceClassification, AlbertTokenizer
 import numpy as np
 import os
 import gdown
@@ -10,31 +10,49 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Define Google Drive folder IDs for each model (use specific subfolder IDs)
-model_drive_ids = {
-    "sentiment": "1uHY8dme-adxXsq7KrqoHjT6jhCtHZ4xc",
-    "emotion": "1pHCJ2eqd9hHlfqNrRagV0sEszYwwQY2a",
-    "hate_speech": "1th6peD5GBtdSVdW9pPKAPRFn_I12RNiz",
-    "sarcasm": "1gjvxD7WoJx0V7AqtWPNFU_c4NmeFTRO8"
+# Define Google Drive file IDs for each model's config and safetensors
+model_file_ids = {
+    "sentiment": {
+        "config": "11jwMJmQMGkiVZWBRQ5BLFyot1520FYIQ",
+        "model": "115N5yiu9lfw4uJE5YxHNoHauHeYSSusu"
+    },
+    "emotion": {
+        "config": "1dSxK10jbZyRpMDCm6MCRf9Jy0weOzLP9",
+        "model": "1Y3rTtPfo4zu28OhsRybdJF6czZN46I0Y"
+    },
+    "hate_speech": {
+        "config": "1QTejES8BZQs3qnxom9ymiZkLRUAZ91NP",
+        "model": "1ol2xO4XbdHwP_HHCYsnX8iVutA6javy_"
+    },
+    "sarcasm": {
+        "config": "1ypl0j1Yp_-0szR4-P1-0CMyDYBwUn5Wz",
+        "model": "1pbByLvTIHO_sT9HMeypvXbsdHsLVzTdk"
+    }
 }
 
 # 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():
+# Download individual model files
+for task, files in model_file_ids.items():
     output_dir = os.path.join(save_dir, task)
-    if not os.path.exists(output_dir):
-        logger.info(f"Downloading {task} model from Google Drive...")
-        gdown.download_folder(
-            f"https://drive.google.com/drive/folders/1kEXKoJxxD5-0FO8WvtagzseSIC5q-rRY?usp=sharing/{folder_id}",
-            output=output_dir,
-            quiet=False,
-            use_cookies=False
-        )
+    os.makedirs(output_dir, exist_ok=True)
+    
+    config_path = os.path.join(output_dir, "config.json")
+    model_path = os.path.join(output_dir, "model.safetensors")
+    
+    if not os.path.exists(config_path):
+        logger.info(f"Downloading {task} config.json from Google Drive...")
+        gdown.download(f"https://drive.google.com/uc?id={files['config']}", config_path, quiet=False)
+    else:
+        logger.info(f"Config for {task} already exists, skipping download.")
+    
+    if not os.path.exists(model_path):
+        logger.info(f"Downloading {task} model.safetensors from Google Drive...")
+        gdown.download(f"https://drive.google.com/uc?id={files['model']}", model_path, quiet=False)
     else:
-        logger.info(f"Model directory {output_dir} already exists, skipping download.")
+        logger.info(f"Model for {task} already exists, skipping download.")
 
 # Define model paths
 tasks = ["sentiment", "emotion", "hate_speech", "sarcasm"]
@@ -48,10 +66,11 @@ label_mappings = {
     "sarcasm": ["no", "yes"]
 }
 
-# Load tokenizer with use_fast=False to avoid tiktoken dependency
+# Load tokenizer
 logger.info("Loading tokenizer...")
 try:
-    tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indic-bert", use_fast=False)
+    # Explicitly use AlbertTokenizer with SentencePiece
+    tokenizer = AlbertTokenizer.from_pretrained("ai4bharat/indic-bert", use_fast=False)
 except Exception as e:
     logger.error(f"Failed to load tokenizer: {str(e)}")
     raise